embedchain

1
import hashlib
2
from unittest.mock import MagicMock
3

4
import pytest
5

6
from embedchain.loaders.mysql import MySQLLoader
7

8

9
@pytest.fixture
10
def mysql_loader(mocker):
11
    with mocker.patch("mysql.connector.connection.MySQLConnection"):
12
        config = {
13
            "host": "localhost",
14
            "port": "3306",
15
            "user": "your_username",
16
            "password": "your_password",
17
            "database": "your_database",
18
        }
19
        loader = MySQLLoader(config=config)
20
        yield loader
21

22

23
def test_mysql_loader_initialization(mysql_loader):
24
    assert mysql_loader.config is not None
25
    assert mysql_loader.connection is not None
26
    assert mysql_loader.cursor is not None
27

28

29
def test_mysql_loader_invalid_config():
30
    with pytest.raises(ValueError, match="Invalid sql config: None"):
31
        MySQLLoader(config=None)
32

33

34
def test_mysql_loader_setup_loader_successful(mysql_loader):
35
    assert mysql_loader.connection is not None
36
    assert mysql_loader.cursor is not None
37

38

39
def test_mysql_loader_setup_loader_connection_error(mysql_loader, mocker):
40
    mocker.patch("mysql.connector.connection.MySQLConnection", side_effect=IOError("Mocked connection error"))
41
    with pytest.raises(ValueError, match="Unable to connect with the given config:"):
42
        mysql_loader._setup_loader(config={})
43

44

45
def test_mysql_loader_check_query_successful(mysql_loader):
46
    query = "SELECT * FROM table"
47
    mysql_loader._check_query(query=query)
48

49

50
def test_mysql_loader_check_query_invalid(mysql_loader):
51
    with pytest.raises(ValueError, match="Invalid mysql query: 123"):
52
        mysql_loader._check_query(query=123)
53

54

55
def test_mysql_loader_load_data_successful(mysql_loader, mocker):
56
    mock_cursor = MagicMock()
57
    mocker.patch.object(mysql_loader, "cursor", mock_cursor)
58
    mock_cursor.fetchall.return_value = [(1, "data1"), (2, "data2")]
59

60
    query = "SELECT * FROM table"
61
    result = mysql_loader.load_data(query)
62

63
    assert "doc_id" in result
64
    assert "data" in result
65
    assert len(result["data"]) == 2
66
    assert result["data"][0]["meta_data"]["url"] == query
67
    assert result["data"][1]["meta_data"]["url"] == query
68

69
    doc_id = hashlib.sha256((query + ", ".join([d["content"] for d in result["data"]])).encode()).hexdigest()
70

71
    assert result["doc_id"] == doc_id
72
    assert mock_cursor.execute.called_with(query)
73

74

75
def test_mysql_loader_load_data_invalid_query(mysql_loader):
76
    with pytest.raises(ValueError, match="Invalid mysql query: 123"):
77
        mysql_loader.load_data(query=123)
78