fix: Update SSL parameter to boolean in Postgres connection and add identity fields in Azure Search datasource tests

Author: Pavan-Microsoft

code/tests/chat_history/test_postgresdbservice.py

@@ -39,7 +39,7 @@ async def test_connect(mock_credential, mock_connect, postgres_client, mock_conn
         database="test_db",
         password="mock_token",
         port=5432,
-        ssl="require",
+        ssl=True,
     )
     assert postgres_client.conn == mock_connection
 

code/tests/functional/tests/functions/integrated_vectorization/test_integrated_vectorization_resource_creation.py

@@ -102,6 +102,10 @@ def test_integrated_vectorization_datasouce_created(
                 "container": {
                     "name": f"{app_config.get_from_json('AZURE_BLOB_STORAGE_INFO','containerName')}"
                 },
+                "identity": {
+                    "@odata.type": "#Microsoft.Azure.Search.DataUserAssignedIdentity",
+                    "userAssignedIdentity": ""
+                },
                 "dataDeletionDetectionPolicy": {
                     "@odata.type": "#Microsoft.Azure.Search.NativeBlobSoftDeleteDeletionDetectionPolicy"
                 },
@@ -367,6 +371,10 @@ def test_integrated_vectorization_skillset_created(
                         "resourceUri": f"https://localhost:{httpserver.port}/",
                         "deploymentId": f"{app_config.get_from_json('AZURE_OPENAI_EMBEDDING_MODEL_INFO','model')}",
                         "apiKey": f"{app_config.get('AZURE_OPENAI_API_KEY')}",
+                        "authIdentity": {
+                            "@odata.type": "#Microsoft.Azure.Search.DataUserAssignedIdentity",
+                            "userAssignedIdentity": ""
+                        },
                     },
                 ],
                 "indexProjections": {

code/tests/utilities/helpers/test_azure_postgres_helper.py

@@ -30,7 +30,7 @@ def test_create_search_client_success(self, mock_connect, mock_credential):
             "https://ossrdbms-aad.database.windows.net/.default"
         )
         mock_connect.assert_called_once_with(
-            "host=mock_host user=mock_user dbname=mock_database password=mock-access-token"
+            "host=mock_host user=mock_user dbname=mock_database password=mock-access-token sslmode=require"
         )
 
     @patch("backend.batch.utilities.helpers.azure_postgres_helper.psycopg2.connect")
@@ -92,7 +92,7 @@ def test_get_vector_store_success(self, mock_cursor, mock_connect, mock_credenti
         # Assert
         self.assertEqual(results, mock_results)
         mock_connect.assert_called_once_with(
-            "host=mock_host user=mock_user dbname=mock_database password=mock-access-token"
+            "host=mock_host user=mock_user dbname=mock_database password=mock-access-token sslmode=require"
         )
 
     @patch("backend.batch.utilities.helpers.azure_postgres_helper.get_azure_credential")

code/tests/utilities/helpers/test_azure_search_helper.py

@@ -90,6 +90,7 @@ def env_helper_mock():
         env_helper.AZURE_SEARCH_CONVERSATIONS_LOG_INDEX = (
             AZURE_SEARCH_CONVERSATIONS_LOG_INDEX
         )
+        env_helper.MANAGED_IDENTITY_CLIENT_ID = "mock-client-id"
 
         env_helper.USE_ADVANCED_IMAGE_PROCESSING = USE_ADVANCED_IMAGE_PROCESSING
         env_helper.is_auth_type_keys.return_value = True
@@ -156,7 +157,7 @@ def test_creates_search_clients_with_rabc(
     AzureSearchHelper()
 
     # then
-    default_azure_credential_mock.assert_called_once_with()
+    default_azure_credential_mock.assert_called_once_with("mock-client-id")
     search_client_mock.assert_called_once_with(
         endpoint=AZURE_SEARCH_SERVICE,
         index_name=AZURE_SEARCH_INDEX,

code/tests/utilities/helpers/test_secret_helper.py

@@ -25,6 +25,7 @@ def test_get_secret_returns_value_from_secret_client_when_use_key_vault_is_true(
     secret_name = "MY_SECRET"
     expected_value = ""
     monkeypatch.setenv("USE_KEY_VAULT", "true")
+    monkeypatch.setenv("AZURE_KEY_VAULT_ENDPOINT", "https://test-vault.vault.azure.net/")
     secret_client.return_value.get_secret.return_value.value = expected_value
     secret_helper = SecretHelper()
 

code/tests/utilities/integrated_vectorization/test_azure_search_datasource.py

@@ -6,6 +6,9 @@
 from azure.search.documents.indexes._generated.models import (
     NativeBlobSoftDeleteDeletionDetectionPolicy,
 )
+from azure.search.documents.indexes.models import (
+    SearchIndexerDataUserAssignedIdentity,
+)
 
 AZURE_AUTH_TYPE = "keys"
 AZURE_SEARCH_KEY = "mock-key"
@@ -18,6 +21,8 @@
 AZURE_BLOB_ACCOUNT_KEY = "mock-key"
 AZURE_SUBSCRIPTION_ID = "mock-subscriptionid"
 AZURE_RESOURCE_GROUP = "mock-resource-group"
+AZURE_BLOB_CONTAINER_NAME = "mock-container-name"
+MANAGED_IDENTITY_RESOURCE_ID = "/subscriptions/mock-sub/resourceGroups/mock-rg/providers/Microsoft.ManagedIdentity/userAssignedIdentities/mock-identity"
 
 
 @pytest.fixture(autouse=True)
@@ -33,6 +38,13 @@ def env_helper_mock():
         env_helper.AZURE_OPENAI_ENDPOINT = AZURE_OPENAI_ENDPOINT
         env_helper.AZURE_OPENAI_EMBEDDING_MODEL = AZURE_OPENAI_EMBEDDING_MODEL
         env_helper.AZURE_SEARCH_DATASOURCE_NAME = AZURE_SEARCH_DATASOURCE_NAME
+        env_helper.AZURE_BLOB_ACCOUNT_NAME = AZURE_BLOB_ACCOUNT_NAME
+        env_helper.AZURE_BLOB_ACCOUNT_KEY = AZURE_BLOB_ACCOUNT_KEY
+        env_helper.AZURE_SUBSCRIPTION_ID = AZURE_SUBSCRIPTION_ID
+        env_helper.AZURE_RESOURCE_GROUP = AZURE_RESOURCE_GROUP
+        env_helper.AZURE_BLOB_CONTAINER_NAME = AZURE_BLOB_CONTAINER_NAME
+        env_helper.MANAGED_IDENTITY_RESOURCE_ID = MANAGED_IDENTITY_RESOURCE_ID
+        env_helper.APP_ENV = "prod"
 
         yield env_helper
 
@@ -89,6 +101,9 @@ def test_create_or_update_datasource_keys(
         connection_string=keys_datasource_connection,
         container=search_indexer_data_container_mock.return_value,
         data_deletion_detection_policy=NativeBlobSoftDeleteDeletionDetectionPolicy(),
+        identity=SearchIndexerDataUserAssignedIdentity(
+            user_assigned_identity=env_helper_mock.MANAGED_IDENTITY_RESOURCE_ID
+        ),
     )
 
 
@@ -123,4 +138,43 @@ def test_create_or_update_datasource_rbac(
         connection_string=rbac_datasource_connection,
         container=search_indexer_data_container_mock.return_value,
         data_deletion_detection_policy=NativeBlobSoftDeleteDeletionDetectionPolicy(),
+        identity=SearchIndexerDataUserAssignedIdentity(
+            user_assigned_identity=env_helper_mock.MANAGED_IDENTITY_RESOURCE_ID
+        ),
+    )
+
+
+def test_create_or_update_datasource_dev_environment(
+    search_indexer_client_mock: MagicMock,
+    search_indexer_data_container_mock: MagicMock,
+    env_helper_mock: MagicMock,
+    search_indexer_datasource_connection_mock: MagicMock,
+):
+    # given
+    env_helper_mock.is_auth_type_keys.return_value = False
+    env_helper_mock.AZURE_AUTH_TYPE = "rbac"
+    env_helper_mock.APP_ENV = "dev"  # Override for dev environment
+    rbac_datasource_connection = f"ResourceId=/subscriptions/{env_helper_mock.AZURE_SUBSCRIPTION_ID}/resourceGroups/{env_helper_mock.AZURE_RESOURCE_GROUP}/providers/Microsoft.Storage/storageAccounts/{env_helper_mock.AZURE_BLOB_ACCOUNT_NAME}/;"
+
+    azure_search_iv_datasource_helper = AzureSearchDatasource(env_helper_mock)
+
+    # when
+    azure_search_iv_datasource_helper.create_or_update_datasource()
+
+    # then
+
+    assert (
+        azure_search_iv_datasource_helper.indexer_client
+        == search_indexer_client_mock.return_value
+    )
+    search_indexer_data_container_mock.assert_called_once_with(
+        name=env_helper_mock.AZURE_BLOB_CONTAINER_NAME
+    )
+    search_indexer_datasource_connection_mock.assert_called_once_with(
+        name=env_helper_mock.AZURE_SEARCH_DATASOURCE_NAME,
+        type="azureblob",
+        connection_string=rbac_datasource_connection,
+        container=search_indexer_data_container_mock.return_value,
+        data_deletion_detection_policy=NativeBlobSoftDeleteDeletionDetectionPolicy(),
+        identity=None,
     )