Coverage for mindsdb / integrations / handlers / mongodb_handler / tests / test_mongodb_handler.py: 0%

1import pytest 

2import json 

3from pymongo import MongoClient 

4 

5from mindsdb_sql_parser.ast import Identifier, Select, Star 

6from mindsdb.integrations.handlers.mongodb_handler.mongodb_handler import MongoDBHandler 

7from mindsdb.integrations.libs.response import RESPONSE_TYPE 

8 

9 

10HANDLER_KWARGS = { 

11 "connection_data": { 

12 "host": "127.0.0.1", 

13 "port": "27017", 

14 "username": "test_user", 

15 "password": "supersecret", 

16 "database": "mongo_test_db", 

17 } 

18} 

19 

20expected_columns = ["_id", "col_one", "col_two", "col_three", "col_four"] 

21 

22 

23def seed_db(): 

24 """Seed the test DB with some data""" 

25 creds = HANDLER_KWARGS["connection_data"] 

26 uri = f"mongodb://{creds['username']}:{creds['password']}@{creds['host']}" 

27 conn = MongoClient(uri) 

28 db = conn[HANDLER_KWARGS["connection_data"]["database"]] # noqa 

29 

30 with open("mindsdb/integrations/handlers/mongodb_handler/tests/seed.json", "r") as f: 

31 json.load(f) 

32 conn.close() 

33 

34 

35@pytest.fixture(scope="module") 

36def handler(request): 

37 seed_db() 

38 handler = MongoDBHandler("mongo_handler", **HANDLER_KWARGS) 

39 return handler 

40 

41 

42def check_valid_response(res): 

43 if res.resp_type == RESPONSE_TYPE.TABLE: 

44 assert res.data_frame is not None, "expected to have some data, but got None" 

45 assert ( 

46 res.error_code == 0 

47 ), f"expected to have zero error_code, but got {res.error_code}" 

48 assert ( 

49 res.error_message is None 

50 ), f"expected to have None in error message, but got {res.error_message}" 

51 

52 

53""" TESTS """ 

54 

55# TODO - Subscribe 

56 

57 

58class TestMongoDBConnection: 

59 def test_connect(self, handler): 

60 handler.connect() 

61 assert handler.is_connected, "connection error" 

62 

63 def test_check_connection(self, handler): 

64 res = handler.check_connection() 

65 assert res.success is True, res.error_message 

66 

67 

68# TODO - Subscribe 

69 

70 

71class TestMongoDBQuery: 

72 def test_native_query(self, handler): 

73 query_string = "db.test.find()" 

74 response = handler.native_query(query_string) 

75 dbs = response.data_frame 

76 assert dbs is not None, "expected to get some data, but got None" 

77 assert "col_one" in dbs, f"expected to get 'col_one' column in response:\n{dbs}" 

78 

79 def test_select_query(self, handler): 

80 limit = 3 

81 query = Select( 

82 targets=[Star()], 

83 from_table=Identifier(parts=["test"]), 

84 ) 

85 res = handler.query(query) 

86 check_valid_response(res) 

87 got_rows = res.data_frame.shape[0] 

88 want_rows = limit 

89 assert ( 

90 got_rows == want_rows 

91 ), f"expected to have {want_rows} rows in response but got: {got_rows}" 

92 

93 

94class TestMongoDBTables: 

95 def test_get_tables(self, handler): 

96 res = handler.get_tables() 

97 tables = res.data_frame 

98 test_table = list(tables["table_name"]) 

99 assert ( 

100 tables is not None 

101 ), "expected to have some table_name in the db, but got None" 

102 assert ( 

103 "table_name" in tables 

104 ), f"expected to get 'table_name' column in the response:\n{tables}" 

105 # get a specific table from the tables list 

106 assert ( 

107 "test" in test_table 

108 ), f"expected to have 'test' table in the db but got: {test_table}" 

109 

110 

111class TestMongoDBColumns: 

112 def test_get_columns(self, handler): 

113 described = handler.get_columns("test") 

114 describe_data = described.data_frame 

115 check_valid_response(described) 

116 got_columns = list(describe_data.iloc[:, 0]) 

117 assert ( 

118 got_columns == expected_columns 

119 ), f"expected to have next columns in test table:\n{expected_columns}\nbut got:\n{got_columns}" 

120 

121 

122class TestMongoDBDisconnect: 

123 def test_disconnect(self, handler): 

124 handler.disconnect() 

125 assert handler.is_connected is False, "failed to disconnect"