Coverage for crateweb/nlp_classification/tests/database_connection_tests.py: 100%

1from unittest import mock 

2from django.test import TestCase 

3 

4from crate_anon.crateweb.nlp_classification.database_connection import ( 

5 DatabaseConnection, 

6) 

7 

8 

9class DatabaseConnectionTest(TestCase): 

10 def test_fetchone_as_dict_no_condition(self) -> None: 

11 column_names = ["column_one", "column_two", "column_three"] 

12 table_name = "test_table" 

13 

14 mock_execute = mock.Mock() 

15 mock_fetchone = mock.Mock(return_value=[1, 2, 3]) 

16 mock_cursor = mock.Mock(execute=mock_execute, fetchone=mock_fetchone) 

17 mock_cm = mock.Mock( 

18 return_value=mock.Mock( 

19 __enter__=mock.Mock(return_value=mock_cursor), 

20 __exit__=mock.Mock(), 

21 ) 

22 ) 

23 connection = DatabaseConnection("test") 

24 mock_connections = {"test": mock.Mock(cursor=mock_cm)} 

25 

26 with mock.patch.multiple( 

27 "crate_anon.crateweb.nlp_classification.database_connection", 

28 connections=mock_connections, 

29 ): 

30 row_dict = connection.fetchone_as_dict(column_names, table_name) 

31 self.assertEqual( 

32 row_dict, {"column_one": 1, "column_two": 2, "column_three": 3} 

33 ) 

34 

35 mock_execute.assert_called_once_with( 

36 "SELECT column_one, column_two, column_three FROM test_table", 

37 None, 

38 ) 

39 mock_fetchone.assert_called_once_with() 

40 

41 def test_fetchone_as_dict_with_condition(self) -> None: 

42 column_names = ["column_one", "column_two", "column_three"] 

43 table_name = "test_table" 

44 

45 mock_execute = mock.Mock() 

46 mock_fetchone = mock.Mock(return_value=[1, 2, 3]) 

47 mock_cursor = mock.Mock(execute=mock_execute, fetchone=mock_fetchone) 

48 mock_cm = mock.Mock( 

49 return_value=mock.Mock( 

50 __enter__=mock.Mock(return_value=mock_cursor), 

51 __exit__=mock.Mock(), 

52 ) 

53 ) 

54 connection = DatabaseConnection("test") 

55 mock_connections = {"test": mock.Mock(cursor=mock_cm)} 

56 

57 with mock.patch.multiple( 

58 "crate_anon.crateweb.nlp_classification.database_connection", 

59 connections=mock_connections, 

60 ): 

61 connection.fetchone_as_dict( 

62 column_names, table_name, "column_one = %s", ["1"] 

63 ) 

64 mock_execute.assert_called_once_with( 

65 ( 

66 "SELECT column_one, column_two, column_three " 

67 "FROM test_table " 

68 "WHERE column_one = %s" 

69 ), 

70 ["1"], 

71 )