Coverage for mindsdb / integrations / handlers / lancedb_handler / tests / test_lancedb_handler.py: 0%

1import unittest 

2from mindsdb.api.executor.data_types.response_type import ( 

3 RESPONSE_TYPE, 

4) 

5from mindsdb.integrations.handlers.lancedb_handler.lancedb_handler import ( 

6 LanceDBHandler, 

7) 

8from mindsdb.integrations.libs.vectordatabase_handler import ( 

9 FilterCondition, 

10 FilterOperator, 

11) 

12import pandas as pd 

13 

14 

15class LanceDBHandlerTest(unittest.TestCase): 

16 @classmethod 

17 def setUpClass(cls): 

18 cls.kwargs = {'connection_data': {'persist_directory': '~/lancedb'}} 

19 cls.handler = LanceDBHandler('test_lancedb_handler', **cls.kwargs) 

20 

21 def test_0_connect(self): 

22 self.handler.connect() 

23 

24 def test_1_check_connection(self): 

25 self.handler.check_connection() 

26 

27 def test_2_create_table(self): 

28 res = self.handler.create_table('test_data17') 

29 assert res.resp_type is not RESPONSE_TYPE.ERROR 

30 

31 def test_3_insert_into_table(self): 

32 data = [{'id': 'id1', 'content': 'this is a test', 'metadata': {'test': 'test1'}, 'embeddings': [1, 2, 3, 4, 3, 5, 2, 8]}, 

33 {'id': 'id2', 'content': 'this is a test', 'metadata': {'test': 'test2'}, 'embeddings': [4, 2, 7, 4, 2, 5, 2, 9]}, 

34 {'id': 'id3', 'content': 'this is a test', 'metadata': {'test': 'test3'}, 'embeddings': [5, 2, 3, 2, 3, 3, 2, 7]}, 

35 {'id': 'id3', 'content': 'this is a test', 'metadata': {'test': 'test4'}, 'embeddings': [5, 2, 3, 2, 3, 4, 2, 7]}] 

36 df = pd.DataFrame(data) 

37 res = self.handler.insert('test_data17', df, None) 

38 assert res.resp_type is not RESPONSE_TYPE.ERROR 

39 

40 def test_4_select(self): 

41 res = self.handler.select( 

42 'test_data17', 

43 ['id', 'content', 'metadata', 'embeddings'], 

44 [ 

45 FilterCondition( 

46 column="id", 

47 op=FilterOperator.EQUAL, 

48 value="id3", 

49 ) 

50 ], 

51 None, 

52 None, 

53 ) 

54 assert res.resp_type is RESPONSE_TYPE.TABLE 

55 

56 def test_5_vector_distance(self): 

57 res = self.handler.select( 

58 'test_data17', 

59 ['id', 'content', 'metadata', 'embeddings'], 

60 [ 

61 FilterCondition( 

62 column="search_vector", 

63 op=FilterOperator.EQUAL, 

64 value="[4.0, 2.0, 7.0, 4.0, 2.0, 5.0, 2.0, 9.0]", 

65 ) 

66 ], 

67 None, 

68 None, 

69 ) 

70 assert res.resp_type is RESPONSE_TYPE.TABLE and 'distance' in res.data_frame.columns 

71 

72 def test_6_delete(self): 

73 res = self.handler.delete( 

74 'test_data17', 

75 [ 

76 FilterCondition( 

77 column="id", 

78 op=FilterOperator.EQUAL, 

79 value="id1", 

80 ) 

81 ] 

82 ) 

83 assert res.resp_type is not RESPONSE_TYPE.ERROR 

84 

85 def test_7_describe_table(self): 

86 res = self.handler.get_columns('test_data17') 

87 assert res.resp_type is RESPONSE_TYPE.TABLE 

88 

89 def test_8_get_tables(self): 

90 res = self.handler.get_tables() 

91 assert res.resp_type is not RESPONSE_TYPE.ERROR 

92 

93 def test_9_drop_table(self): 

94 res = self.handler.drop_table('test_data17') 

95 assert res.resp_type is not RESPONSE_TYPE.ERROR 

96 

97 

98if __name__ == '__main__': 

99 unittest.main()