rename AzureServicePrincipalMigration to ServicePrincipalMigration, r…

…ename execute_migration to run

src/databricks/labs/ucx/azure/azure_credentials.py

@@ -54,7 +54,7 @@ def from_storage_credential_validation(
         )
 
 
-class AzureServicePrincipalMigration:
+class ServicePrincipalMigration:
 
     def __init__(
         self,
@@ -65,7 +65,6 @@ def __init__(
         integration_test_flag="",
     ):
         self._output_file = "azure_service_principal_migration_result.csv"
-        self._final_sp_list: list[ServicePrincipalMigrationInfo] = []
         self._installation = installation
         self._ws = ws
         self._azure_resource_permissions = azure_resource_permissions
@@ -111,7 +110,7 @@ def _list_storage_credentials(self) -> set[str]:
                     # return the storage credential created during integration test
                     return {storage_credential.azure_service_principal.application_id}
             # return no storage credential if there is none created during integration test
-            return {""}
+            return {}
 
         for storage_credential in storage_credentials:
             # only add service principal's application_id, ignore managed identity based storage_credential
@@ -186,17 +185,17 @@ def _fetch_client_secret(self, sp_list: list[StoragePermissionMapping]) -> list[
                 )
         return sp_list_with_secret
 
-    def _print_action_plan(self, sp_list_with_secret: list[ServicePrincipalMigrationInfo]):
+    def _print_action_plan(self, sp_list: list[ServicePrincipalMigrationInfo]):
         # print action plan to console for customer to review.
-        for sp in sp_list_with_secret:
-            print(
+        for sp in sp_list:
+            logger.info(
                 f"Service Principal name: {sp.service_principal.principal}, "
                 f"application_id: {sp.service_principal.client_id}, "
                 f"privilege {sp.service_principal.privilege} "
                 f"on location {sp.service_principal.prefix}"
             )
 
-    def _generate_migration_list(self):
+    def _generate_migration_list(self) -> list[ServicePrincipalMigrationInfo]:
         """
         Create the list of SP that need to be migrated, output an action plan as a csv file for users to confirm
         """
@@ -208,9 +207,9 @@ def _generate_migration_list(self):
         filtered_sp_list = [sp for sp in sp_list if sp.client_id not in sc_set]
         # fetch sp client_secret if any
         sp_list_with_secret = self._fetch_client_secret(filtered_sp_list)
-        self._final_sp_list = sp_list_with_secret
         # output the action plan for customer to confirm
         self._print_action_plan(sp_list_with_secret)
+        return sp_list_with_secret
 
     def _create_storage_credential(self, sp_migration: ServicePrincipalMigrationInfo):
         # prepare the storage credential properties
@@ -260,9 +259,9 @@ def _validate_storage_credential(
                 ),
             )
 
-    def execute_migration(self, prompts: Prompts):
+    def run(self, prompts: Prompts):
 
-        self._generate_migration_list()
+        sp_list_with_secret = self._generate_migration_list()
 
         plan_confirmed = prompts.confirm(
             "Above Azure Service Principals will be migrated to UC storage credentials, please review and confirm."
@@ -271,7 +270,7 @@ def execute_migration(self, prompts: Prompts):
             return
 
         execution_result = []
-        for sp in self._final_sp_list:
+        for sp in sp_list_with_secret:
             execution_result.append(self._create_storage_credential(sp))
 
         results_file = self._installation.save(execution_result, filename=self._output_file)

src/databricks/labs/ucx/cli.py

@@ -13,7 +13,7 @@
 from databricks.labs.ucx.account import AccountWorkspaces, WorkspaceInfo
 from databricks.labs.ucx.assessment.aws import AWSResourcePermissions
 from databricks.labs.ucx.azure.access import AzureResourcePermissions
-from databricks.labs.ucx.azure.azure_credentials import AzureServicePrincipalMigration
+from databricks.labs.ucx.azure.azure_credentials import ServicePrincipalMigration
 from databricks.labs.ucx.config import WorkspaceConfig
 from databricks.labs.ucx.framework.crawlers import StatementExecutionBackend
 from databricks.labs.ucx.hive_metastore import ExternalLocations, TablesCrawler
@@ -297,8 +297,8 @@ def migrate_azure_service_principals(w: WorkspaceClient):
     """
     logger.info("Running migrate_azure_service_principals command")
     prompts = Prompts()
-    service_principal_migration = AzureServicePrincipalMigration.for_cli(w, prompts)
-    service_principal_migration.execute_migration(prompts)
+    service_principal_migration = ServicePrincipalMigration.for_cli(w, prompts)
+    service_principal_migration.run(prompts)
 
 
 if __name__ == "__main__":

tests/integration/azure/test_azure_credentials.py

@@ -1,14 +1,14 @@
 import base64
 import re
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import pytest
 from databricks.labs.blueprint.tui import MockPrompts
 
 from databricks.labs.ucx.assessment.azure import AzureServicePrincipalInfo
 from databricks.labs.ucx.assessment.crawlers import _SECRET_PATTERN
 from databricks.labs.ucx.azure.access import StoragePermissionMapping
-from databricks.labs.ucx.azure.azure_credentials import AzureServicePrincipalMigration
+from databricks.labs.ucx.azure.azure_credentials import ServicePrincipalMigration
 
 
 @pytest.fixture
@@ -75,15 +75,15 @@ def inner(read_only=False) -> dict:
 
 @pytest.fixture
 def execute_migration(ws):
-    def inner(variables: dict, integration_test_flag: str) -> AzureServicePrincipalMigration:
-        spn_migration = AzureServicePrincipalMigration(
+    def inner(variables: dict, integration_test_flag: str) -> ServicePrincipalMigration:
+        spn_migration = ServicePrincipalMigration(
             variables["installation"],
             ws,
             variables["azure_resource_permissions"],
             variables["azure_sp_crawler"],
             integration_test_flag=integration_test_flag,
         )
-        spn_migration.execute_migration(
+        spn_migration.run(
             MockPrompts({"Above Azure Service Principals will be migrated to UC storage credentials *": "Yes"})
         )
         return spn_migration
@@ -92,7 +92,7 @@ def inner(variables: dict, integration_test_flag: str) -> AzureServicePrincipalM
 
 
 def test_spn_migration_existed_storage_credential(
-    ws, execute_migration, make_storage_credential_from_spn, prepare_spn_migration_test
+    execute_migration, make_storage_credential_from_spn, prepare_spn_migration_test
 ):
     variables = prepare_spn_migration_test(read_only=False)
 
@@ -104,11 +104,11 @@ def test_spn_migration_existed_storage_credential(
         directory_id=variables["directory_id"],
     )
 
-    # test that the spn migration will be skipped due to above storage credential is existed
-    spn_migration = execute_migration(variables, integration_test_flag=variables["storage_credential_name"])
-
-    # because storage_credential is existing, no spn should be migrated
-    assert not spn_migration._final_sp_list
+    with patch("databricks.labs.ucx.azure.azure_credentials.ServicePrincipalMigration._create_storage_credential") as create_storage_credential:
+        # test that the spn migration will be skipped due to above storage credential is existed
+        execute_migration(variables, integration_test_flag=variables["storage_credential_name"])
+        # because storage_credential is existing, no spn should be migrated
+        create_storage_credential.assert_not_called()
 
 
 @pytest.mark.parametrize("read_only", [False, True])
@@ -118,8 +118,9 @@ def test_spn_migration(ws, execute_migration, prepare_spn_migration_test, read_o
     try:
         spn_migration = execute_migration(variables, integration_test_flag="lets_migrate_the_spn")
 
-        assert spn_migration._final_sp_list[0].service_principal.principal == variables["storage_credential_name"]
-        assert ws.storage_credentials.get(variables["storage_credential_name"]).read_only is read_only
+        storage_credential = ws.storage_credentials.get(variables["storage_credential_name"])
+        assert storage_credential is not None
+        assert storage_credential.read_only is read_only
 
         validation_result = spn_migration._installation.save.call_args.args[0][0]
         if read_only:

tests/unit/azure/test_azure_credentials.py

@@ -25,7 +25,7 @@
 )
 from databricks.labs.ucx.azure.access import StoragePermissionMapping
 from databricks.labs.ucx.azure.azure_credentials import (
-    AzureServicePrincipalMigration,
+    ServicePrincipalMigration,
     ServicePrincipalMigrationInfo,
 )
 from tests.unit.framework.mocks import MockBackend
@@ -103,21 +103,21 @@ def download(path: str) -> io.BytesIO:
 
 def test_for_cli_not_azure(caplog, ws):
     ws.config.is_azure = False
-    assert AzureServicePrincipalMigration.for_cli(ws, MagicMock()) is None
+    assert ServicePrincipalMigration.for_cli(ws, MagicMock()) is None
     assert "Workspace is not on azure, please run this command on azure databricks workspaces." in caplog.text
 
 
 def test_for_cli_not_prompts(ws):
     ws.config.is_azure = True
     prompts = MockPrompts({"Have you reviewed the azure_storage_account_info.csv *": "No"})
-    assert AzureServicePrincipalMigration.for_cli(ws, prompts) is None
+    assert ServicePrincipalMigration.for_cli(ws, prompts) is None
 
 
 def test_for_cli(ws):
     ws.config.is_azure = True
     prompts = MockPrompts({"Have you reviewed the azure_storage_account_info.csv *": "Yes"})
 
-    assert isinstance(AzureServicePrincipalMigration.for_cli(ws, prompts), AzureServicePrincipalMigration)
+    assert isinstance(ServicePrincipalMigration.for_cli(ws, prompts), ServicePrincipalMigration)
 
 
 def test_list_storage_credentials(ws):
@@ -137,7 +137,7 @@ def test_list_storage_credentials(ws):
         ),
     ]
 
-    sp_migration = AzureServicePrincipalMigration(MagicMock(), ws, MagicMock(), MagicMock())
+    sp_migration = ServicePrincipalMigration(MagicMock(), ws, MagicMock(), MagicMock())
 
     expected = {"b6420590-5e1c-4426-8950-a94cbe9b6115"}
     sp_migration._list_storage_credentials()
@@ -165,7 +165,7 @@ def test_list_storage_credentials_for_integration_test(ws):
 
     # test storage credential: spn_for_integration_test is picked up
     # in integration test, we only pick up the existing storage credential created in integration test and ignore the others
-    sp_migration = AzureServicePrincipalMigration(
+    sp_migration = ServicePrincipalMigration(
         MagicMock(), ws, MagicMock(), MagicMock(), integration_test_flag="spn_for_integration_test"
     )
     expected = {"b6420590-5e1c-4426-8950-a94cbe9b6115"}
@@ -174,11 +174,11 @@ def test_list_storage_credentials_for_integration_test(ws):
 
     # test storage credential is not picked up
     # if integration test does not create storage credential, we use dummy integration_test_flag to filter out other existing storage credentials
-    sp_migration = AzureServicePrincipalMigration(
+    sp_migration = ServicePrincipalMigration(
         MagicMock(), ws, MagicMock(), MagicMock(), integration_test_flag="other_spn"
     )
     sp_migration._list_storage_credentials()
-    assert {""} == sp_migration._list_storage_credentials()
+    assert {} == sp_migration._list_storage_credentials()
 
 
 @pytest.mark.parametrize(
@@ -192,7 +192,7 @@ def test_list_storage_credentials_for_integration_test(ws):
 def test_read_secret_value_decode(ws, secret_bytes_value, expected_return):
     ws.secrets.get_secret.return_value = secret_bytes_value
 
-    sp_migration = AzureServicePrincipalMigration(MagicMock(), ws, MagicMock(), MagicMock())
+    sp_migration = ServicePrincipalMigration(MagicMock(), ws, MagicMock(), MagicMock())
     assert sp_migration._read_databricks_secret("test_scope", "test_key", "000") == expected_return
 
 
@@ -207,7 +207,7 @@ def test_read_secret_read_exception(caplog, ws, exception, expected_log, expecte
     caplog.set_level(logging.INFO)
     ws.secrets.get_secret.side_effect = exception
 
-    sp_migration = AzureServicePrincipalMigration(MagicMock(), ws, MagicMock(), MagicMock())
+    sp_migration = ServicePrincipalMigration(MagicMock(), ws, MagicMock(), MagicMock())
     secret_value = sp_migration._read_databricks_secret("test_scope", "test_key", "000")
 
     assert expected_log in caplog.text
@@ -281,13 +281,14 @@ def test_fetch_client_secret(ws):
         ),
     ]
 
-    sp_migration = AzureServicePrincipalMigration(MagicMock(), ws, MagicMock(), sp_crawler)
+    sp_migration = ServicePrincipalMigration(MagicMock(), ws, 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):
+def test_print_action_plan(caplog):
+    caplog.set_level(logging.INFO)
     sp_list_with_secret = [
         ServicePrincipalMigrationInfo(
             StoragePermissionMapping(
@@ -300,19 +301,20 @@ def test_print_action_plan(capsys):
             "hello world",
         )
     ]
-    sp_migration = AzureServicePrincipalMigration(MagicMock(), MagicMock(), MagicMock(), MagicMock())
+    sp_migration = ServicePrincipalMigration(MagicMock(), MagicMock(), MagicMock(), MagicMock())
     sp_migration._print_action_plan(sp_list_with_secret)
 
     expected_print = (
         "Service Principal name: principal_1, "
         "application_id: app_secret1, "
         "privilege WRITE_FILES "
-        "on location prefix1\n"
+        "on location prefix1"
     )
-    assert expected_print == capsys.readouterr().out
+    assert expected_print in caplog.text
 
 
-def test_generate_migration_list(capsys, mocker, ws):
+def test_generate_migration_list(caplog, mocker, ws):
+    caplog.set_level(logging.INFO)
     ws.config.is_azure = True
     ws.secrets.get_secret.return_value = GetSecretResponse(value="aGVsbG8gd29ybGQ=")
     ws.storage_credentials.list.return_value = [
@@ -335,10 +337,11 @@ def test_generate_migration_list(capsys, mocker, ws):
         ],
     )
 
-    sp_migration = AzureServicePrincipalMigration.for_cli(ws, prompts)
+    sp_migration = ServicePrincipalMigration.for_cli(ws, prompts)
     sp_migration._generate_migration_list()
 
-    assert "app_secret2" in capsys.readouterr().out
+    assert "app_secret2" in caplog.text
+    assert "app_secret1" not in caplog.text
 
 
 def test_execute_migration_no_confirmation(mocker, ws):
@@ -350,13 +353,13 @@ def test_execute_migration_no_confirmation(mocker, ws):
         }
     )
 
-    mocker.patch("databricks.labs.ucx.azure.azure_credentials.AzureServicePrincipalMigration._generate_migration_list")
+    mocker.patch("databricks.labs.ucx.azure.azure_credentials.ServicePrincipalMigration._generate_migration_list")
 
     with patch(
-        "databricks.labs.ucx.azure.azure_credentials.AzureServicePrincipalMigration._create_storage_credential"
+        "databricks.labs.ucx.azure.azure_credentials.ServicePrincipalMigration._create_storage_credential"
     ) as c:
-        sp_migration = AzureServicePrincipalMigration.for_cli(ws, prompts)
-        sp_migration.execute_migration(prompts)
+        sp_migration = ServicePrincipalMigration.for_cli(ws, prompts)
+        sp_migration.run(prompts)
         c.assert_not_called()
 
 
@@ -409,9 +412,9 @@ def test_execute_migration(caplog, capsys, mocker, ws):
         ],
     )
 
-    sp_migration = AzureServicePrincipalMigration.for_cli(ws, prompts)
+    sp_migration = ServicePrincipalMigration.for_cli(ws, prompts)
     sp_migration._installation.save = MagicMock()
-    sp_migration.execute_migration(prompts)
+    sp_migration.run(prompts)
 
     # assert migration is complete
     assert "Completed migration" in capsys.readouterr().out