Skip to content

Commit

Permalink
add unit test for _generate_migration_list
Browse files Browse the repository at this point in the history
  • Loading branch information
qziyuan committed Feb 21, 2024
1 parent a2a51cf commit 1fa336c
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 3 deletions.
3 changes: 1 addition & 2 deletions src/databricks/labs/ucx/migration/azure_credentials.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'):
Expand Down Expand Up @@ -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
Expand Down
25 changes: 24 additions & 1 deletion tests/unit/migration/test_azure_credentials.py
Original file line number Diff line number Diff line change
@@ -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 (
Expand Down Expand Up @@ -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

0 comments on commit 1fa336c

Please sign in to comment.