Coverage for mindsdb / integrations / handlers / ckan_handler / tests / test_ckan_handler.py: 0%

1import unittest 

2from unittest.mock import patch, MagicMock 

3from mindsdb.integrations.handlers.ckan_handler.ckan_handler import CkanHandler 

4from mindsdb.api.executor.data_types.response_type import RESPONSE_TYPE 

5from mindsdb_sql_parser import ast 

6 

7 

8class CkanHandlerTest(unittest.TestCase): 

9 def setUp(self): 

10 self.handler = CkanHandler('test_ckan_handler', connection_data={ 

11 "url": "http://mock-ckan-url.com", 

12 "api_key": "mock_api_key" 

13 }) 

14 self.patcher = patch('mindsdb.integrations.handlers.ckan_handler.ckan_handler.RemoteCKAN') 

15 self.mock_ckan = self.patcher.start() 

16 self.mock_ckan_instance = MagicMock() 

17 self.mock_ckan.return_value = self.mock_ckan_instance 

18 

19 def tearDown(self): 

20 self.patcher.stop() 

21 

22 def test_check_connection(self): 

23 self.mock_ckan_instance.action.site_read.return_value = True 

24 response = self.handler.check_connection() 

25 self.assertTrue(response.success) 

26 

27 def test_get_tables(self): 

28 self.mock_ckan_instance.action.package_list.return_value = [ 

29 "package1", 

30 "package2", 

31 ] 

32 result = self.handler.get_tables() 

33 self.assertIsNotNone(result) 

34 self.assertEqual(result.type, RESPONSE_TYPE.TABLE) 

35 

36 def test_query_package_ids(self): 

37 mock_packages = [ 

38 { 

39 "id": "pkg1", 

40 "name": "Package 1", 

41 "title": "Test Package 1", 

42 "num_resources": 2, 

43 }, 

44 { 

45 "id": "pkg2", 

46 "name": "Package 2", 

47 "title": "Test Package 2", 

48 "num_resources": 1, 

49 }, 

50 ] 

51 self.mock_ckan_instance.action.package_search.return_value = { 

52 "results": mock_packages 

53 } 

54 

55 query = ast.Select( 

56 targets=[ast.Star()], 

57 from_table=ast.Identifier("package_ids"), 

58 limit=ast.Constant(10), 

59 ) 

60 result = self.handler.query(query) 

61 self.assertEqual(result.type, RESPONSE_TYPE.TABLE) 

62 self.assertEqual(len(result.data_frame), 2) 

63 

64 def test_query_resource_ids(self): 

65 mock_packages = [ 

66 { 

67 "id": "pkg1", 

68 "resources": [ 

69 { 

70 "id": "res1", 

71 "name": "Resource 1", 

72 "format": "CSV", 

73 "datastore_active": True, 

74 }, 

75 { 

76 "id": "res2", 

77 "name": "Resource 2", 

78 "format": "XSLX", 

79 "datastore_active": False, 

80 }, 

81 ], 

82 }, 

83 { 

84 "id": "pkg2", 

85 "resources": [ 

86 { 

87 "id": "res3", 

88 "name": "Resource 3", 

89 "format": "CSV", 

90 "datastore_active": True, 

91 } 

92 ], 

93 }, 

94 ] 

95 self.mock_ckan_instance.action.package_search.return_value = { 

96 "results": mock_packages 

97 } 

98 

99 query = ast.Select( 

100 targets=[ast.Star()], 

101 from_table=ast.Identifier("resource_ids"), 

102 limit=ast.Constant(10), 

103 ) 

104 result = self.handler.query(query) 

105 self.assertEqual(result.type, RESPONSE_TYPE.TABLE) 

106 self.assertEqual(len(result.data_frame), 2) 

107 

108 def test_query_datastore_without_resource_id(self): 

109 query = ast.Select( 

110 targets=[ast.Star()], 

111 from_table=ast.Identifier("datastore"), 

112 limit=ast.Constant(10), 

113 ) 

114 result = self.handler.query(query) 

115 self.assertEqual(result.type, RESPONSE_TYPE.TABLE) 

116 self.assertEqual(len(result.data_frame), 1) 

117 self.assertIn("message", result.data_frame.columns) 

118 

119 def test_query_datastore_with_resource_id(self): 

120 mock_records = [ 

121 {"id": 1, "name": "Record 1"}, 

122 {"id": 2, "name": "Record 2"}, 

123 {"id": 3, "name": "Record 3"}, 

124 ] 

125 self.mock_ckan_instance.action.datastore_search.return_value = { 

126 "records": mock_records 

127 } 

128 

129 query = ast.Select( 

130 targets=[ast.Star()], 

131 from_table=ast.Identifier("datastore"), 

132 where=ast.BinaryOperation( 

133 "=", 

134 args=[ast.Identifier("resource_id"), ast.Constant("test_resource_id")], 

135 ), 

136 limit=ast.Constant(10), 

137 ) 

138 result = self.handler.query(query) 

139 self.assertEqual(result.type, RESPONSE_TYPE.TABLE) 

140 self.assertEqual(len(result.data_frame), 3) 

141 

142 def test_native_query(self): 

143 self.mock_ckan_instance.action.package_list.return_value = [ 

144 "package1", 

145 "package2", 

146 "package3", 

147 ] 

148 

149 query = "package_list:" 

150 result = self.handler.native_query(query) 

151 self.assertEqual(result.type, RESPONSE_TYPE.TABLE) 

152 self.assertEqual(len(result.data_frame), 3) 

153 

154 

155if __name__ == "__main__": 

156 unittest.main()