From 428b8d02ce0a49d6bd81139ab17a3fa980ba1bca Mon Sep 17 00:00:00 2001 From: Anton Shchederkin Date: Mon, 16 Dec 2024 18:11:23 +0100 Subject: [PATCH] Fix CI --- .../src/databricks_cdk/resources/handler.py | 10 ++++-- .../service_principals/service_principal.py | 10 ++++-- aws-lambda/tests/conftest.py | 2 ++ .../test_service_principal.py | 34 +++++++++---------- 4 files changed, 34 insertions(+), 22 deletions(-) diff --git a/aws-lambda/src/databricks_cdk/resources/handler.py b/aws-lambda/src/databricks_cdk/resources/handler.py index b6d890c3..eaa59c5b 100644 --- a/aws-lambda/src/databricks_cdk/resources/handler.py +++ b/aws-lambda/src/databricks_cdk/resources/handler.py @@ -91,6 +91,11 @@ create_or_update_secret_scope, delete_secret_scope, ) +from databricks_cdk.resources.service_principals.service_principal import ( + ServicePrincipalProperties, + create_or_update_service_principal, + delete_service_principal, +) from databricks_cdk.resources.sql_warehouses.sql_warehouses import ( SQLWarehouseProperties, create_or_update_warehouse, @@ -125,7 +130,6 @@ delete_storage_credential, ) from databricks_cdk.resources.unity_catalog.volumes import VolumeProperties, create_or_update_volume, delete_volume -from databricks_cdk.resources.service_principals.service_principal import create_or_update_service_principal, delete_service_principal, ServicePrincipalProperties from databricks_cdk.utils import CnfResponse logger = logging.getLogger(__name__) @@ -329,7 +333,9 @@ def delete_resource(event: DatabricksEvent) -> CnfResponse: elif action == "volume": return delete_volume(VolumeProperties(**event.ResourceProperties), event.PhysicalResourceId) elif action == "service-principal": - return delete_service_principal(ServicePrincipalProperties(**event.ResourceProperties), event.PhysicalResourceId) + return delete_service_principal( + ServicePrincipalProperties(**event.ResourceProperties), event.PhysicalResourceId + ) else: raise RuntimeError(f"Unknown action: {action}") diff --git a/aws-lambda/src/databricks_cdk/resources/service_principals/service_principal.py b/aws-lambda/src/databricks_cdk/resources/service_principals/service_principal.py index cc3672cb..9dc8ae4e 100644 --- a/aws-lambda/src/databricks_cdk/resources/service_principals/service_principal.py +++ b/aws-lambda/src/databricks_cdk/resources/service_principals/service_principal.py @@ -52,11 +52,13 @@ def get_service_principal(physical_resource_id: str, workspace_client: Workspace service_principal = workspace_client.service_principals.get(id=physical_resource_id) except NotFound: raise ServicePrincipalNotFoundError(f"Service principal with id {physical_resource_id} not found") - + return service_principal -def create_service_principal(service_principal: ServicePrincipal, workspace_client: WorkspaceClient) -> ServicePrincipalResponse: +def create_service_principal( + service_principal: ServicePrincipal, workspace_client: WorkspaceClient +) -> ServicePrincipalResponse: """Create service principal on databricks""" created_service_principal = workspace_client.service_principals.create( @@ -74,7 +76,9 @@ def create_service_principal(service_principal: ServicePrincipal, workspace_clie if created_service_principal.id is None: raise ServicePrincipalCreationError("Service principal creation failed, there was no id found") - return ServicePrincipalResponse(name=created_service_principal.display_name, physical_resource_id=created_service_principal.id) + return ServicePrincipalResponse( + name=created_service_principal.display_name, physical_resource_id=created_service_principal.id + ) def update_service_principal( diff --git a/aws-lambda/tests/conftest.py b/aws-lambda/tests/conftest.py index 297f4f8d..9cde9d33 100644 --- a/aws-lambda/tests/conftest.py +++ b/aws-lambda/tests/conftest.py @@ -3,6 +3,7 @@ import pytest from databricks.sdk import AccountClient, CredentialsAPI, ExperimentsAPI, ModelRegistryAPI, VolumesAPI, WorkspaceClient +from databricks.sdk.service.iam import ServicePrincipalsAPI @pytest.fixture(scope="function", autouse=True) @@ -24,6 +25,7 @@ def workspace_client(): workspace_client.model_registry = MagicMock(spec=ModelRegistryAPI) workspace_client.experiments = MagicMock(spec=ExperimentsAPI) workspace_client.volumes = MagicMock(spec=VolumesAPI) + workspace_client.service_principals = MagicMock(spec=ServicePrincipalsAPI) return workspace_client diff --git a/aws-lambda/tests/resources/service_principals/test_service_principal.py b/aws-lambda/tests/resources/service_principals/test_service_principal.py index af8a025a..9c1501b1 100644 --- a/aws-lambda/tests/resources/service_principals/test_service_principal.py +++ b/aws-lambda/tests/resources/service_principals/test_service_principal.py @@ -1,19 +1,19 @@ from unittest.mock import patch import pytest - -from databricks.sdk.service.iam import ServicePrincipal, ComplexValue from databricks.sdk.errors import NotFound +from databricks.sdk.service.iam import ComplexValue, ServicePrincipal + from databricks_cdk.resources.service_principals.service_principal import ( ServicePrincipalCreationError, + ServicePrincipalNotFoundError, ServicePrincipalProperties, ServicePrincipalResponse, - ServicePrincipalNotFoundError, create_or_update_service_principal, create_service_principal, delete_service_principal, - update_service_principal, get_service_principal, + update_service_principal, ) from databricks_cdk.utils import CnfResponse @@ -24,10 +24,10 @@ @patch("databricks_cdk.resources.service_principals.service_principal.get_service_principal") def test_create_or_update_service_principal_create( patched_get_service_principal, - patched_create_service_principal, - patched_update_service_principal, - patched_get_workspace_client, - workspace_client + patched_create_service_principal, + patched_update_service_principal, + patched_get_workspace_client, + workspace_client, ): patched_get_workspace_client.return_value = workspace_client mock_properties = ServicePrincipalProperties( @@ -42,7 +42,7 @@ def test_create_or_update_service_principal_create( patched_create_service_principal.assert_called_once_with( ServicePrincipal( - active=True, + active=True, display_name="mock_name", id=None, application_id=None, @@ -51,8 +51,8 @@ def test_create_or_update_service_principal_create( groups=None, roles=None, schemas=None, - ), - workspace_client + ), + workspace_client, ) patched_get_service_principal.assert_not_called() patched_update_service_principal.assert_not_called() @@ -64,10 +64,10 @@ def test_create_or_update_service_principal_create( @patch("databricks_cdk.resources.service_principals.service_principal.get_service_principal") def test_create_or_update_service_principal_update( patched_get_service_principal, - patched_create_service_principal, - patched_update_service_principal, - patched_get_workspace_client, - workspace_client + patched_create_service_principal, + patched_update_service_principal, + patched_get_workspace_client, + workspace_client, ): patched_get_workspace_client.return_value = workspace_client mock_physical_resource_id = "some_id" @@ -101,8 +101,8 @@ def test_create_or_update_service_principal_update( external_id=None, groups=None, schemas=None, - ), - workspace_client + ), + workspace_client, ) patched_create_service_principal.assert_not_called()