mock default credentials

Fr4nc3 · Dec 2, 2024 · 6198de7 · 6198de7
1 parent 270386c
commit 6198de7
Showing 1 changed file with 29 additions and 3 deletions.
diff --git a/code/tests/utilities/helpers/test_azure_postgres_helper.py b/code/tests/utilities/helpers/test_azure_postgres_helper.py
@@ -101,8 +101,11 @@ def test_get_search_indexes_success(
             "host=mock_host user=mock_user dbname=mock_database password=mock-access-token"
         )
 
+    @patch(
+        "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential"
+    )
     @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect")
-    def test_get_search_indexes_query_error(self, mock_connect):
+    def test_get_search_indexes_query_error(self, mock_connect, mock_credential):
         # Arrange
         # Mock the EnvHelper and set required attributes
         mock_env_helper = MagicMock()
@@ -111,6 +114,11 @@ def test_get_search_indexes_query_error(self, mock_connect):
         mock_env_helper.POSTGRESQL_DATABASE = "mock_database"
         mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5
 
+        # Mock access token retrieval
+        mock_access_token = MagicMock()
+        mock_access_token.token = "mock-access-token"
+        mock_credential.return_value.get_token.return_value = mock_access_token
+
         mock_connection = MagicMock()
         mock_connect.return_value = mock_connection
 
@@ -815,18 +823,26 @@ def test_search_by_blob_url_success(
         mock_connection.close.assert_called_once()
         mock_logger.info.assert_called_with("Retrieved 2 unique title(s).")
 
+    @patch(
+        "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential"
+    )
     @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect")
     @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger")
     @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper")
     def test_search_by_blob_url_no_results(
-        self, mock_env_helper, mock_logger, mock_connect
+        self, mock_env_helper, mock_logger, mock_connect, mock_credential
     ):
         # Arrange: Mock the EnvHelper attributes
         mock_env_helper.POSTGRESQL_USER = "mock_user"
         mock_env_helper.POSTGRESQL_HOST = "mock_host"
         mock_env_helper.POSTGRESQL_DATABASE = "mock_database"
         mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5
 
+        # Mock access token retrieval
+        mock_access_token = MagicMock()
+        mock_access_token.token = "mock-access-token"
+        mock_credential.return_value.get_token.return_value = mock_access_token
+
         # Mock the connection and cursor
         mock_connection = MagicMock()
         mock_cursor = MagicMock()
@@ -850,16 +866,26 @@ def test_search_by_blob_url_no_results(
         mock_connection.close.assert_called_once()
         mock_logger.info.assert_called_with("Retrieved 0 unique title(s).")
 
+    @patch(
+        "backend.batch.utilities.helpers.azure_postgres_helper.DefaultAzureCredential"
+    )
     @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect")
     @patch("backend.batch.utilities.helpers.azure_postgres_helper.logger")
     @patch("backend.batch.utilities.helpers.azure_postgres_helper.EnvHelper")
-    def test_search_by_blob_url_error(self, mock_env_helper, mock_logger, mock_connect):
+    def test_search_by_blob_url_error(
+        self, mock_env_helper, mock_logger, mock_connect, mock_credential
+    ):
         # Arrange: Mock the EnvHelper attributes
         mock_env_helper.POSTGRESQL_USER = "mock_user"
         mock_env_helper.POSTGRESQL_HOST = "mock_host"
         mock_env_helper.POSTGRESQL_DATABASE = "mock_database"
         mock_env_helper.AZURE_POSTGRES_SEARCH_TOP_K = 5
 
+        # Mock access token retrieval
+        mock_access_token = MagicMock()
+        mock_access_token.token = "mock-access-token"
+        mock_credential.return_value.get_token.return_value = mock_access_token
+
         # Mock the connection and cursor
         mock_connection = MagicMock()
         mock_cursor = MagicMock()