fix leftover and failed unit tests after rebase

databrickslabs · Feb 21, 2024 · e7a3867 · e7a3867
1 parent b55a812
commit e7a3867
Show file tree

Hide file tree

Showing 6 changed files with 9 additions and 12 deletions.
diff --git a/src/databricks/labs/ucx/azure/access.py b/src/databricks/labs/ucx/azure/access.py
@@ -65,7 +65,7 @@ def _map_storage(self, storage: AzureResource) -> list[StoragePermissionMapping]
                         client_id=role_assignment.principal.client_id,
                         principal=role_assignment.principal.display_name,
                         privilege=privilege,
-                        directory_id = role_assignment.principal.directory_id
+                        directory_id=role_assignment.principal.directory_id,
                     )
                 )
         return out

diff --git a/src/databricks/labs/ucx/azure/azure_credentials.py b/src/databricks/labs/ucx/azure/azure_credentials.py
@@ -16,12 +16,12 @@
     ValidationResult,
 )
 
-from databricks.labs.ucx.assessment.azure import (
+from databricks.labs.ucx.assessment.azure import AzureServicePrincipalCrawler
+from databricks.labs.ucx.azure.access import (
     AzureResourcePermissions,
-    AzureResources,
-    AzureServicePrincipalCrawler,
     StoragePermissionMapping,
 )
+from databricks.labs.ucx.azure.resources import AzureResources
 from databricks.labs.ucx.config import WorkspaceConfig
 from databricks.labs.ucx.framework.crawlers import StatementExecutionBackend
 from databricks.labs.ucx.hive_metastore.locations import ExternalLocations

diff --git a/src/databricks/labs/ucx/cli.py b/src/databricks/labs/ucx/cli.py
@@ -12,7 +12,6 @@
 
 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.config import WorkspaceConfig

diff --git a/tests/integration/azure/test_azure_credentials.py b/tests/integration/azure/test_azure_credentials.py
@@ -5,11 +5,9 @@
 import pytest
 from databricks.labs.blueprint.tui import MockPrompts
 
-from databricks.labs.ucx.assessment.azure import (
-    AzureServicePrincipalInfo,
-    StoragePermissionMapping,
-)
+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
 
 

diff --git a/tests/unit/azure/test_azure_credentials.py b/tests/unit/azure/test_azure_credentials.py
@@ -22,8 +22,8 @@
 from databricks.labs.ucx.assessment.azure import (
     AzureServicePrincipalCrawler,
     AzureServicePrincipalInfo,
-    StoragePermissionMapping,
 )
+from databricks.labs.ucx.azure.access import StoragePermissionMapping
 from databricks.labs.ucx.azure.azure_credentials import (
     AzureServicePrincipalMigration,
     ServicePrincipalMigrationInfo,

diff --git a/tests/unit/azure/test_resources.py b/tests/unit/azure/test_resources.py
@@ -61,7 +61,7 @@ def test_role_assignments_storage(mocker, az_token):
     assert len(role_assignments) == 1
     for role_assignment in role_assignments:
         assert role_assignment.role_name == "Contributor"
-        assert role_assignment.principal == Principal("appIduser2", "disNameuser2", "Iduser2")
+        assert role_assignment.principal == Principal("appIduser2", "disNameuser2", "Iduser2", "0000-0000")
         assert str(role_assignment.scope) == resource_id
         assert role_assignment.resource == AzureResource(resource_id)
 
@@ -75,6 +75,6 @@ def test_role_assignments_container(mocker, az_token):
     assert len(role_assignments) == 1
     for role_assignment in role_assignments:
         assert role_assignment.role_name == "Contributor"
-        assert role_assignment.principal == Principal("appIduser2", "disNameuser2", "Iduser2")
+        assert role_assignment.principal == Principal("appIduser2", "disNameuser2", "Iduser2", "0000-0000")
         assert str(role_assignment.scope) == resource_id
         assert role_assignment.resource == AzureResource(resource_id)