From 1fa336c9cba74a6a1b37abbcea5b8f5006a453ae Mon Sep 17 00:00:00 2001
From: Ziyuan Qin <ziyuan.qin@databricks.com>
Date: Sun, 11 Feb 2024 21:19:41 -0800
Subject: [PATCH] add unit test for _generate_migration_list

---
 .../labs/ucx/migration/azure_credentials.py   |  3 +--
 .../unit/migration/test_azure_credentials.py  | 25 ++++++++++++++++++-
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/src/databricks/labs/ucx/migration/azure_credentials.py b/src/databricks/labs/ucx/migration/azure_credentials.py
index 1f88c48605..1fde03af5f 100644
--- a/src/databricks/labs/ucx/migration/azure_credentials.py
+++ b/src/databricks/labs/ucx/migration/azure_credentials.py
@@ -76,7 +76,6 @@ def __init__(
         self._ws = ws
         self._azure_resource_permissions = azure_resource_permissions
         self._azure_sp_crawler = azure_sp_crawler
-        self._action_plan = 'service_principals_for_storage_credentials.csv'
 
     @classmethod
     def for_cli(cls, ws: WorkspaceClient, prompts: Prompts, product='ucx'):
@@ -202,7 +201,7 @@ def _generate_migration_list(self):
         # list existed storage credentials
         sc_set = self._list_storage_credentials()
         # check if the sp is already used in UC storage credential
-        filtered_sp_list = [sp for sp in sp_list if sp not in sc_set]
+        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
diff --git a/tests/unit/migration/test_azure_credentials.py b/tests/unit/migration/test_azure_credentials.py
index 4e9a336dfc..7e76de970b 100644
--- 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, Mock
+from unittest.mock import MagicMock, create_autospec, Mock, patch
 
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import (
@@ -139,5 +139,28 @@ def test_print_action_plan(capsys):
     assert expected_print == capsys.readouterr().out
 
 
+def test_generate_migration_list(capsys, mocker, ws):
+    ws.config.is_azure.return_value = True
+    ws.secrets.get_secret.return_value = GetSecretResponse(value="aGVsbG8gd29ybGQ=")
+    ws.storage_credentials.list.return_value = [
+        StorageCredentialInfo(
+            azure_service_principal=AzureServicePrincipal(
+                application_id="app_secret1",
+                directory_id="directory_id_1",
+                client_secret="hello world",
+            )
+        )
+    ]
+
+    prompts = MockPrompts({"Have you reviewed the azure_storage_account_info.csv *": "Yes"})
+
+    mocker.patch("databricks.labs.ucx.assessment.azure.AzureResourcePermissions.load", return_value = [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")])
+    mocker.patch("databricks.labs.ucx.assessment.azure.AzureServicePrincipalCrawler.snapshot", return_value=[AzureServicePrincipalInfo("app_secret1", "test_scope", "test_key", "tenant_id_1", "storage1"),
+                                                                                                             AzureServicePrincipalInfo("app_secret2", "test_scope", "test_key", "tenant_id_1", "storage1")])
+
+    sp_migration = AzureServicePrincipalMigration.for_cli(ws, prompts)
+    sp_migration._generate_migration_list()
 
+    assert "app_secret2" in capsys.readouterr().out