dd unit tests for _fetch_client_secret and _print_action_plan

qziyuan · qziyuan · commit b365488093a4 · 2024-02-21T14:18:16.000-08:00
diff --git a/src/databricks/labs/ucx/migration/azure_credentials.py b/src/databricks/labs/ucx/migration/azure_credentials.py
@@ -163,9 +163,9 @@ def _fetch_client_secret(self, sp_list: list[StoragePermissionMapping]) -> list[
         azure_sp_infos = self._azure_sp_crawler.snapshot()
 
         for azure_sp_info in azure_sp_infos:
-            if azure_sp_info.secret_scope is None:
+            if not azure_sp_info.secret_scope:
                 continue
-            if azure_sp_info.secret_key is None:
+            if not azure_sp_info.secret_key:
                 continue
             secret_value = self._read_databricks_secret(
                 azure_sp_info.secret_scope, azure_sp_info.secret_key, azure_sp_info.application_id
diff --git a/tests/unit/migration/test_azure_credentials.py b/tests/unit/migration/test_azure_credentials.py
@@ -1,7 +1,7 @@
 import logging
 import pytest
 
-from unittest.mock import MagicMock, create_autospec
+from unittest.mock import MagicMock, create_autospec, Mock
 
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import (
@@ -16,10 +16,12 @@
 )
 from databricks.sdk.service.workspace import GetSecretResponse
 
-from databricks.labs.ucx.assessment.azure import StoragePermissionMapping
+from databricks.labs.ucx.assessment.azure import StoragePermissionMapping, \
+    AzureServicePrincipalCrawler, AzureServicePrincipalInfo
 from databricks.labs.ucx.migration.azure_credentials import (
-    AzureServicePrincipalMigration,
+    AzureServicePrincipalMigration, ServicePrincipalMigrationInfo,
 )
+from tests.unit.framework.mocks import MockBackend
 
 
 def test_list_storage_credentials():
@@ -74,4 +76,46 @@ def test_read_secret_read_exception(caplog, exception, expected_log, expected_re
     secret_value = sp_migration._read_databricks_secret("test_scope","test_key", "000")
 
     assert expected_log in caplog.text
-    assert secret_value == expected_return
+    assert secret_value == expected_return
+
+
+def test_fetch_client_secret():
+    w = create_autospec(WorkspaceClient)
+    w.secrets.get_secret.return_value = GetSecretResponse(value="aGVsbG8gd29ybGQ=")
+
+    crawled_sp = [AzureServicePrincipalInfo("app_secret1", "test_scope", "test_key", "tenant_id_1", "storage1"),
+                  AzureServicePrincipalInfo("app_secret2", "test_scope", "test_key", "tenant_id_1", "storage1"),
+                  AzureServicePrincipalInfo("app_no_secret1", "", "", "tenant_id_1", "storage1"),
+                  AzureServicePrincipalInfo("app_no_secret2", "test_scope", "", "tenant_id_1", "storage1"),]
+    sp_crawler = AzureServicePrincipalCrawler(w, MockBackend(), "ucx")
+    sp_crawler._try_fetch = Mock(return_value=crawled_sp)
+    sp_crawler._crawl = Mock(return_value=crawled_sp)
+
+    sp_to_be_checked = [StoragePermissionMapping(prefix="prefix1",client_id="app_secret1",principal="principal_1",privilege="WRITE_FILES",directory_id="directory_id_1"),
+                       StoragePermissionMapping(prefix="prefix2",client_id="app_secret2",principal="principal_2",privilege="READ_FILES",directory_id="directory_id_1"),
+                       StoragePermissionMapping(prefix="prefix3",client_id="app_no_secret1",principal="principal_3",privilege="WRITE_FILES",directory_id="directory_id_2"),
+                       StoragePermissionMapping(prefix="prefix4",client_id="app_no_secret2",principal="principal_4",privilege="READ_FILES",directory_id="directory_id_2")]
+
+    expected_sp_list = [ServicePrincipalMigrationInfo(StoragePermissionMapping(prefix="prefix1",client_id="app_secret1",principal="principal_1",privilege="WRITE_FILES",directory_id="directory_id_1"), "hello world"),
+                        ServicePrincipalMigrationInfo(StoragePermissionMapping(prefix="prefix2",client_id="app_secret2",principal="principal_2",privilege="READ_FILES",directory_id="directory_id_1"), "hello world")]
+
+    sp_migration = AzureServicePrincipalMigration(MagicMock(), w, MagicMock(), sp_crawler)
+    filtered_sp_list = sp_migration._fetch_client_secret(sp_to_be_checked)
+
+    assert filtered_sp_list == expected_sp_list
+
+
+def test_print_action_plan(capsys):
+    sp_list_with_secret = [ServicePrincipalMigrationInfo(StoragePermissionMapping(prefix="prefix1",client_id="app_secret1",principal="principal_1",privilege="WRITE_FILES",directory_id="directory_id_1"), "hello world")]
+    sp_migration = AzureServicePrincipalMigration(MagicMock(), MagicMock(), MagicMock(), MagicMock())
+    sp_migration._print_action_plan(sp_list_with_secret)
+
+    expected_print = (f"Service Principal name: principal_1, "
+                      f"application_id: app_secret1, "
+                      f"privilege WRITE_FILES "
+                      f"on location prefix1\n")
+    assert expected_print == capsys.readouterr().out
+
+
+
+
