Coverage for mindsdb / integrations / handlers / athena_handler / tests / test_athena_handler.py: 0%

1import unittest 

2from unittest.mock import patch, MagicMock, Mock 

3from collections import OrderedDict 

4from mindsdb.integrations.handlers.athena_handler.athena_handler import AthenaHandler 

5 

6 

7class CursorContextManager(Mock): 

8 def __enter__(self): 

9 return self 

10 

11 def __exit__(self, *args): 

12 pass 

13 

14 description = [['a']] 

15 

16 def fetchall(self): 

17 return [[1]] 

18 

19 

20class AthenaHandlerTest(unittest.TestCase): 

21 dummy_connection_data = OrderedDict( 

22 aws_access_key_id='aws_access_key_id', 

23 aws_secret_access_key='aws_secret_access_key', 

24 region_name='us-east-1', 

25 database='default', 

26 workgroup='my_workgroup', 

27 catalog='AwsDataCatalog', 

28 results_output_location='s3://bucket-path/athena-query-results', 

29 check_interval=0 

30 ) 

31 

32 def setUp(self): 

33 self.patcher = patch('boto3.client') 

34 self.mock_client = self.patcher.start() 

35 self.mock_client.return_value = MagicMock() 

36 self.handler = AthenaHandler('athena', connection_data=self.dummy_connection_data) 

37 

38 def tearDown(self): 

39 self.patcher.stop() 

40 

41 def test_connect_success(self): 

42 connection = self.handler.connect() 

43 self.assertIsNotNone(connection) 

44 self.assertTrue(self.handler.is_connected) 

45 

46 def test_get_columns(self): 

47 self.handler.native_query = MagicMock() 

48 

49 table_name = "mock_table" 

50 self.handler.get_columns(table_name) 

51 

52 expected_query = f""" 

53 select 

54 column_name as "Field", 

55 data_type as "Type" 

56 from 

57 information_schema.columns 

58 where 

59 table_name = '{table_name}' 

60 """ 

61 

62 self.handler.native_query.assert_called_once_with(expected_query) 

63 

64 def test_get_tables(self): 

65 self.handler.native_query = MagicMock() 

66 

67 self.handler.get_tables() 

68 

69 expected_query = """ 

70 select 

71 table_schema, 

72 table_name, 

73 table_type 

74 from 

75 information_schema.tables 

76 where 

77 table_schema not in ('information_schema') 

78 and table_type in ('BASE TABLE', 'VIEW') 

79 """ 

80 

81 self.handler.native_query.assert_called_once_with(expected_query) 

82 

83 

84if __name__ == '__main__': 

85 unittest.main()