frogml 1.2.38__py3-none-any.whl → 1.2.47__py3-none-any.whl

frogml_services_mock/mocks/environment_v0_service.py

@@ -0,0 +1,175 @@
+from typing import Optional
+
+import grpc
+
+from frogml._proto.qwak.administration.v0.environments.environment_pb2 import (
+    QwakEnvironment,
+)
+from frogml._proto.qwak.administration.v0.environments.environment_service_pb2 import (
+    CreateEnvironmentResponse,
+    GetEnvironmentApplicationUserCredentialsResponse,
+    GetEnvironmentResponse,
+    ListEnvironmentsResponse,
+    SetEnvironmentApplicationUserCredentialsResponse,
+    UpdateEnvironmentConfigurationResponse,
+    UpdateEnvironmentPersonalizationResponse,
+)
+from frogml._proto.qwak.administration.v0.environments.environment_service_pb2_grpc import (
+    EnvironmentServiceServicer,
+)
+
+
+class EnvironmentV0ServiceMock(EnvironmentServiceServicer):
+    def __init__(self):
+        super().__init__()
+        self.__create_environment_response: Optional[CreateEnvironmentResponse] = None
+        self.__create_environment_error: Optional[grpc.StatusCode] = None
+
+        self.__get_environment_response: Optional[GetEnvironmentResponse] = None
+        self.__get_environment_error: Optional[grpc.StatusCode] = None
+
+        self.__update_environment_configuration_response: Optional[
+            UpdateEnvironmentConfigurationResponse
+        ] = None
+        self.__update_environment_configuration_error: Optional[grpc.StatusCode] = None
+
+        self.__update_environment_personalization_response: Optional[
+            UpdateEnvironmentPersonalizationResponse
+        ] = None
+        self.__update_environment_personalization_error: Optional[grpc.StatusCode] = (
+            None
+        )
+
+        self.__list_environments_response: Optional[ListEnvironmentsResponse] = None
+        self.__list_environments_error: Optional[grpc.StatusCode] = None
+
+        self.__get_credentials_response: Optional[
+            GetEnvironmentApplicationUserCredentialsResponse
+        ] = None
+        self.__get_credentials_error: Optional[grpc.StatusCode] = None
+
+        self.__set_credentials_response: Optional[
+            SetEnvironmentApplicationUserCredentialsResponse
+        ] = None
+        self.__set_credentials_error: Optional[grpc.StatusCode] = None
+
+    def given_create_environment(
+        self,
+        response: Optional[CreateEnvironmentResponse] = None,
+        error_code: Optional[grpc.StatusCode] = None,
+    ):
+        self.__create_environment_response = response
+        self.__create_environment_error = error_code
+
+    def given_get_environment(
+        self,
+        response: Optional[GetEnvironmentResponse] = None,
+        error_code: Optional[grpc.StatusCode] = None,
+    ):
+        self.__get_environment_response = response
+        self.__get_environment_error = error_code
+
+    def given_update_environment_configuration(
+        self,
+        response: Optional[UpdateEnvironmentConfigurationResponse] = None,
+        error_code: Optional[grpc.StatusCode] = None,
+    ):
+        self.__update_environment_configuration_response = response
+        self.__update_environment_configuration_error = error_code
+
+    def given_update_environment_personalization(
+        self,
+        response: Optional[UpdateEnvironmentPersonalizationResponse] = None,
+        error_code: Optional[grpc.StatusCode] = None,
+    ):
+        self.__update_environment_personalization_response = response
+        self.__update_environment_personalization_error = error_code
+
+    def given_list_environments(
+        self,
+        response: Optional[ListEnvironmentsResponse] = None,
+        error_code: Optional[grpc.StatusCode] = None,
+    ):
+        self.__list_environments_response = response
+        self.__list_environments_error = error_code
+
+    def given_get_credentials(
+        self,
+        response: Optional[GetEnvironmentApplicationUserCredentialsResponse] = None,
+        error_code: Optional[grpc.StatusCode] = None,
+    ):
+        self.__get_credentials_response = response
+        self.__get_credentials_error = error_code
+
+    def given_set_credentials(
+        self,
+        response: Optional[SetEnvironmentApplicationUserCredentialsResponse] = None,
+        error_code: Optional[grpc.StatusCode] = None,
+    ):
+        self.__set_credentials_response = response
+        self.__set_credentials_error = error_code
+
+    def CreateEnvironment(self, request, context):
+        if self.__create_environment_error:
+            context.set_code(self.__create_environment_error)
+            context.set_details("Failed to create environment")
+            return CreateEnvironmentResponse()
+        if self.__create_environment_response:
+            return self.__create_environment_response
+        return CreateEnvironmentResponse(environment=QwakEnvironment(id="mock-env-id"))
+
+    def GetEnvironment(self, request, context):
+        if self.__get_environment_error:
+            context.set_code(self.__get_environment_error)
+            context.set_details("Failed to get environment")
+            return GetEnvironmentResponse()
+        if self.__get_environment_response:
+            return self.__get_environment_response
+        return GetEnvironmentResponse(
+            environment=QwakEnvironment(id=request.environment_id)
+        )
+
+    def UpdateEnvironmentConfiguration(self, request, context):
+        if self.__update_environment_configuration_error:
+            context.set_code(self.__update_environment_configuration_error)
+            context.set_details("Failed to update environment configuration")
+            return UpdateEnvironmentConfigurationResponse()
+        if self.__update_environment_configuration_response:
+            return self.__update_environment_configuration_response
+        return UpdateEnvironmentConfigurationResponse()
+
+    def UpdateEnvironmentPersonalization(self, request, context):
+        if self.__update_environment_personalization_error:
+            context.set_code(self.__update_environment_personalization_error)
+            context.set_details("Failed to update environment personalization")
+            return UpdateEnvironmentPersonalizationResponse()
+        if self.__update_environment_personalization_response:
+            return self.__update_environment_personalization_response
+        return UpdateEnvironmentPersonalizationResponse()
+
+    def ListEnvironments(self, request, context):
+        if self.__list_environments_error:
+            context.set_code(self.__list_environments_error)
+            context.set_details("Failed to list environments")
+            return ListEnvironmentsResponse()
+        if self.__list_environments_response:
+            return self.__list_environments_response
+        return ListEnvironmentsResponse()
+
+    def GetEnvironmentApplicationUserCredentials(self, request, context):
+        if self.__get_credentials_error:
+            context.set_code(self.__get_credentials_error)
+            context.set_details("Failed to get credentials")
+            return GetEnvironmentApplicationUserCredentialsResponse()
+        if self.__get_credentials_response:
+            return self.__get_credentials_response
+        return GetEnvironmentApplicationUserCredentialsResponse()
+
+    def SetEnvironmentApplicationUserCredentials(self, request, context):
+        if self.__set_credentials_error:
+            context.set_code(self.__set_credentials_error)
+            context.set_details("Failed to set credentials")
+            return SetEnvironmentApplicationUserCredentialsResponse()
+        if self.__set_credentials_response:
+            return self.__set_credentials_response
+        return SetEnvironmentApplicationUserCredentialsResponse()

frogml_services_mock/mocks/environment_v1_service.py

@@ -0,0 +1,37 @@
+from typing import Optional
+
+import grpc
+
+from frogml._proto.qwak.administration.v0.environments.environment_pb2 import (
+    QwakEnvironment,
+)
+from frogml._proto.qwak.administration.v1.environments.environment_service_pb2 import (
+    CreateEnvironmentResponse,
+)
+from frogml._proto.qwak.administration.v1.environments.environment_service_pb2_grpc import (
+    EnvironmentServiceServicer,
+)
+
+
+class EnvironmentV1ServiceMock(EnvironmentServiceServicer):
+    def __init__(self):
+        super().__init__()
+        self.__create_environment_response: Optional[CreateEnvironmentResponse] = None
+        self.__create_environment_error: Optional[grpc.StatusCode] = None
+
+    def given_create_environment(
+        self,
+        response: Optional[CreateEnvironmentResponse] = None,
+        error_code: Optional[grpc.StatusCode] = None,
+    ):
+        self.__create_environment_response = response
+        self.__create_environment_error = error_code
+
+    def CreateEnvironment(self, request, context):
+        if self.__create_environment_error:
+            context.set_code(self.__create_environment_error)
+            context.set_details("Failed to create environment")
+            return CreateEnvironmentResponse()
+        if self.__create_environment_response:
+            return self.__create_environment_response
+        return CreateEnvironmentResponse(environment=QwakEnvironment(id="mock-env-id"))

frogml_services_mock/mocks/frogml_mocks.py

@@ -23,11 +23,14 @@ from frogml_services_mock.mocks.build_orchestrator_build_settings_api import (
 from frogml_services_mock.mocks.build_orchestrator_service_api import (
     BuildOrchestratorServiceApiMock,
 )
+from frogml_services_mock.mocks.cluster_v2_service import ClusterV2ServiceMock
 from frogml_services_mock.mocks.data_versioning_service import DataVersioningServiceMock
 from frogml_services_mock.mocks.deployment_management_service import (
     DeploymentManagementServiceMock,
 )
 from frogml_services_mock.mocks.ecosystem_service_api import EcoSystemServiceMock
+from frogml_services_mock.mocks.environment_v0_service import EnvironmentV0ServiceMock
+from frogml_services_mock.mocks.environment_v1_service import EnvironmentV1ServiceMock
 from frogml_services_mock.mocks.execution_management_service import (
     ExecutionManagementServiceMock,
 )
@@ -56,7 +59,9 @@ from frogml_services_mock.mocks.integration_management_service import (
 from frogml_services_mock.mocks.internal_build_orchestrator_service import (
     InternalBuildOrchestratorServiceMock,
 )
-from frogml_services_mock.mocks.jfrog_tenant_info_service_mock import JFrogTenantInfoServiceMock
+from frogml_services_mock.mocks.jfrog_tenant_info_service_mock import (
+    JFrogTenantInfoServiceMock,
+)
 from frogml_services_mock.mocks.job_registry_service_api import (
     JobRegistryServiceApiMock,
 )
@@ -67,15 +72,20 @@ from frogml_services_mock.mocks.location_discovery_service_api import (
     LocationDiscoveryServiceApiMock,
 )
 from frogml_services_mock.mocks.logging_service import LoggingServiceApiMock
-from frogml_services_mock.mocks.model_group_management_service import ModelGroupManagementServiceMock
+from frogml_services_mock.mocks.model_deployment_manager_service_mock import (
+    ModelDeploymentManagerMock,
+)
+from frogml_services_mock.mocks.model_group_management_service import (
+    ModelGroupManagementServiceMock,
+)
 from frogml_services_mock.mocks.model_management_service import (
     ModelsManagementServiceMock,
 )
 from frogml_services_mock.mocks.model_version_manager_service import (
     ModelVersionManagerServiceMock,
 )
-from frogml_services_mock.mocks.repository_service_mock import RepositoryServiceMock
 from frogml_services_mock.mocks.project_manager_service import ProjectManagerServiceMock
+from frogml_services_mock.mocks.repository_service_mock import RepositoryServiceMock
 from frogml_services_mock.mocks.secret_service import SecretServiceMock
 from frogml_services_mock.mocks.self_service_user_service import (
     SelfServiceUserServiceMock,
@@ -127,3 +137,7 @@ class FrogmlMocks:
     repository_service: RepositoryServiceMock
     location_discovery_service: LocationDiscoveryServiceApiMock
     jfrog_tenant_info_service: JFrogTenantInfoServiceMock
+    model_deployment_manager_mock: ModelDeploymentManagerMock
+    cluster_v2_service_mock: ClusterV2ServiceMock
+    environment_v0_service_mock: EnvironmentV0ServiceMock
+    environment_v1_service_mock: EnvironmentV1ServiceMock

frogml_services_mock/mocks/model_deployment_manager_service_mock.py

@@ -0,0 +1,255 @@
+import uuid
+from datetime import datetime, timezone
+from typing import Callable, Optional
+
+import grpc
+from frogml._proto.com.jfrog.ml.model.deployment.v1.deployment_pb2 import (
+    Deployment,
+    DeploymentMetadata,
+    ModelDeploymentInformation,
+    ModelDeploymentStatus,
+    ModelDeploymentStatusSuccessful,
+    MultipleEnvironmentDeployment,
+)
+from frogml._proto.com.jfrog.ml.model.deployment.v1.deployment_service_pb2 import (
+    DeployModelRequest,
+    DeployModelResponse,
+    EditModelDeploymentRequest,
+    EditModelDeploymentResponse,
+    GetModelDeploymentRequest,
+    GetModelDeploymentResponse,
+    ListModelDeploymentsRequest,
+    ListModelDeploymentsResponse,
+    UndeployModelRequest,
+    UndeployModelResponse,
+)
+from frogml._proto.com.jfrog.ml.model.deployment.v1.deployment_service_pb2_grpc import (
+    ModelDeploymentServiceServicer,
+)
+from frogml._proto.com.jfrog.ml.model.deployment.v1.model_artifact_identifier_pb2 import (
+    ModelArtifactIdentifier,
+    ModelBasedArtifactIdentifier,
+)
+from frogml._proto.com.jfrog.ml.model.deployment.v1.model_deployment_brief_pb2 import (
+    ModelDeploymentBrief,
+)
+from google.protobuf.timestamp_pb2 import Timestamp
+from pydantic import BaseModel, ConfigDict
+
+
+class DeploymentInformation(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    deployment_id: uuid.UUID
+    model_id: str
+    model_group_id: str
+    path: str
+    environment_id: uuid.UUID
+    deployment_config: Deployment
+    created_at: datetime
+
+
+class ModelDeploymentManagerMock(ModelDeploymentServiceServicer):
+    def __init__(self):
+        self.__deployment_id_to_spec: dict[str, DeploymentInformation] = {}
+
+    @staticmethod
+    def __get_model_identifier_details(
+        model_artifact_identifier: ModelArtifactIdentifier,
+    ) -> tuple[str, str, str]:
+        identifier_type = model_artifact_identifier.WhichOneof("identifier_type")
+
+        if identifier_type == "model_based_artifact_id":
+            artifact = model_artifact_identifier.model_based_artifact_id
+            return (
+                artifact.model_id,
+                artifact.model_group_id,
+                artifact.image_path,
+            )
+        elif identifier_type == "custom_model_artifact_id":
+            artifact = model_artifact_identifier.custom_model_artifact_id
+            return (
+                artifact.model_id,
+                artifact.model_group_id,
+                artifact.build_id,
+            )
+
+        raise Exception("No Model Artifact Identifier found")
+
+    def DeployModel(
+        self, request: DeployModelRequest, context: grpc.ServicerContext
+    ) -> DeployModelResponse:
+        multiple_environment_deployment: MultipleEnvironmentDeployment = (
+            request.multiple_environment_deployment
+        )
+        model_artifact_identifier: ModelArtifactIdentifier = (
+            multiple_environment_deployment.model_artifact_identifier
+        )
+        model_id, model_group_id, path = self.__get_model_identifier_details(
+            model_artifact_identifier
+        )
+
+        self.__deployment_id_to_spec.update(
+            {
+                (deployment_id := str(uuid.uuid4())): DeploymentInformation(
+                    deployment_id=uuid.UUID(deployment_id),
+                    model_id=model_id,
+                    model_group_id=model_group_id,
+                    path=path,
+                    environment_id=uuid.UUID(env_id),
+                    deployment_config=deployment,
+                    created_at=datetime.now(tz=timezone.utc),
+                )
+                for env_id, deployment in multiple_environment_deployment.deployments.items()
+            }
+        )
+
+        return DeployModelResponse()
+
+    def UndeployModel(
+        self, request: UndeployModelRequest, context: grpc.ServicerContext
+    ) -> UndeployModelResponse:
+        deployment_ids = list(request.deployment_ids.deployment_id)
+        for deployment_id in deployment_ids:
+            self.__deployment_id_to_spec.pop(deployment_id)
+
+        return UndeployModelResponse()
+
+    def EditModelDeployment(
+        self, request: EditModelDeploymentRequest, context: grpc.ServicerContext
+    ) -> EditModelDeploymentResponse:
+        deployment_map = dict(request.deployment_id_to_deployment_spec_map.deployments)
+        for deployment_id, new_spec in deployment_map.items():
+            if deployment_id not in self.__deployment_id_to_spec:
+                raise Exception("Deployment id does not exist")
+            self.__deployment_id_to_spec[deployment_id].deployment_config = new_spec
+
+        return EditModelDeploymentResponse()
+
+    @staticmethod
+    def __build_filter_predicates(
+        request: ListModelDeploymentsRequest,
+    ) -> list[Callable[[DeploymentInformation], bool]]:
+        """Build a list of filter predicates from the request."""
+        predicates: list[Callable[[DeploymentInformation], bool]] = []
+
+        if not request.model_deployment_filter.HasField(
+            "simple_model_deployment_filter"
+        ):
+            return predicates
+
+        simple_filter = request.model_deployment_filter.simple_model_deployment_filter
+
+        if simple_filter.HasField("model_identifier_filter"):
+            model_group_id = simple_filter.model_identifier_filter.model_group_id
+            model_id = simple_filter.model_identifier_filter.model_id
+            predicates.append(
+                lambda d, model_group_id=model_group_id, model_id=model_id: d.model_group_id
+                == model_group_id
+                and d.model_id == model_id
+            )
+
+        if simple_filter.HasField("model_group_ids"):
+            allowed_model_group_ids = set(simple_filter.model_group_ids.model_group_id)
+            predicates.append(
+                lambda d, allowed=allowed_model_group_ids: d.model_group_id in allowed
+            )
+
+        if simple_filter.HasField("model_ids"):
+            allowed_model_ids = set(simple_filter.model_ids.model_id)
+            predicates.append(
+                lambda d, allowed=allowed_model_ids: d.model_id in allowed
+            )
+
+        if simple_filter.HasField("environment_ids"):
+            allowed_environment_ids = set(simple_filter.environment_ids.environment_id)
+            predicates.append(
+                lambda d, allowed=allowed_environment_ids: str(d.environment_id)
+                in allowed
+            )
+
+        return predicates
+
+    @staticmethod
+    def _deployment_to_brief(
+        deployment_info: DeploymentInformation,
+    ) -> ModelDeploymentBrief:
+        """Convert a DeploymentInformation to a ModelDeploymentBrief."""
+        timestamp = Timestamp()
+        timestamp.FromDatetime(deployment_info.created_at)
+
+        return ModelDeploymentBrief(
+            status=ModelDeploymentStatus(
+                model_deployment_status_successful=ModelDeploymentStatusSuccessful()
+            ),
+            environment=str(deployment_info.environment_id),
+            created_at=timestamp,
+            deployment_id=str(deployment_info.deployment_id),
+            model_artifact_identifier=ModelArtifactIdentifier(
+                model_based_artifact_id=ModelBasedArtifactIdentifier(
+                    model_id=deployment_info.model_id,
+                    model_group_id=deployment_info.model_group_id,
+                    image_path=deployment_info.path,
+                )
+            ),
+        )
+
+    def ListModelDeployments(
+        self, request: ListModelDeploymentsRequest, context: grpc.ServicerContext
+    ) -> ListModelDeploymentsResponse:
+        predicates: list[Callable[[DeploymentInformation], bool]] = (
+            self.__build_filter_predicates(request)
+        )
+
+        result = [
+            self._deployment_to_brief(deployment)
+            for deployment in self.__deployment_id_to_spec.values()
+            if all(predicate(deployment) for predicate in predicates)
+        ]
+
+        return ListModelDeploymentsResponse(model_deployment_brief=result)
+
+    def GetModelDeployment(
+        self, request: GetModelDeploymentRequest, context: grpc.ServicerContext
+    ) -> GetModelDeploymentResponse:
+        deployment_id = request.deployment_id
+        deployment_info = self.__deployment_id_to_spec.get(deployment_id)
+
+        if deployment_info is None:
+            return GetModelDeploymentResponse()
+
+        timestamp = Timestamp()
+        timestamp.FromDatetime(deployment_info.created_at)
+
+        model_artifact_identifier = ModelArtifactIdentifier(
+            model_based_artifact_id=ModelBasedArtifactIdentifier(
+                model_id=deployment_info.model_id,
+                model_group_id=deployment_info.model_group_id,
+                image_path=deployment_info.path,
+            )
+        )
+
+        model_deployment_info = ModelDeploymentInformation(
+            model_artifact_identifier=model_artifact_identifier,
+            model_deployment_spec=deployment_info.deployment_config,
+            model_deployment_metadata=DeploymentMetadata(
+                deployment_id=str(deployment_info.deployment_id),
+                status=ModelDeploymentStatus(
+                    model_deployment_status_successful=ModelDeploymentStatusSuccessful()
+                ),
+                environment=str(deployment_info.environment_id),
+                created_at=timestamp,
+            ),
+        )
+
+        return GetModelDeploymentResponse(
+            model_deployment_information=model_deployment_info
+        )
+
+    def get_deployment_by_id(
+        self, deployment_id: str
+    ) -> Optional[DeploymentInformation]:
+        return self.__deployment_id_to_spec.get(deployment_id)
+
+    def get_all_deployments(self) -> list[DeploymentInformation]:
+        return list(self.__deployment_id_to_spec.values())

frogml_services_mock/services_mock.py

@@ -3,13 +3,24 @@ from typing import Any, Generator, List, Tuple
 
 import grpc
 import pytest
-
+from frogml._proto.com.jfrog.ml.model.deployment.v1.deployment_service_pb2_grpc import (
+    add_ModelDeploymentServiceServicer_to_server,
+)
 from frogml._proto.jfml.model_version.v1.model_version_manager_service_pb2_grpc import (
     add_ModelVersionManagerServiceServicer_to_server,
 )
 from frogml._proto.qwak.administration.v0.authentication.authentication_service_pb2_grpc import (
     add_AuthenticationServiceServicer_to_server,
 )
+from frogml._proto.qwak.administration.cluster.v2.cluster_service_pb2_grpc import (
+    add_ClusterServiceServicer_to_server,
+)
+from frogml._proto.qwak.administration.v0.environments.environment_service_pb2_grpc import (
+    add_EnvironmentServiceServicer_to_server as add_EnvironmentV0ServiceServicer_to_server,
+)
+from frogml._proto.qwak.administration.v1.environments.environment_service_pb2_grpc import (
+    add_EnvironmentServiceServicer_to_server as add_EnvironmentV1ServiceServicer_to_server,
+)
 from frogml._proto.qwak.admiral.secret.v0.system_secret_service_pb2_grpc import (
     add_SystemSecretServiceServicer_to_server,
 )
@@ -142,11 +153,14 @@ from frogml_services_mock.mocks.build_orchestrator_build_settings_api import (
 from frogml_services_mock.mocks.build_orchestrator_service_api import (
     BuildOrchestratorServiceApiMock,
 )
+from frogml_services_mock.mocks.cluster_v2_service import ClusterV2ServiceMock
 from frogml_services_mock.mocks.data_versioning_service import DataVersioningServiceMock
 from frogml_services_mock.mocks.deployment_management_service import (
     DeploymentManagementServiceMock,
 )
 from frogml_services_mock.mocks.ecosystem_service_api import EcoSystemServiceMock
+from frogml_services_mock.mocks.environment_v0_service import EnvironmentV0ServiceMock
+from frogml_services_mock.mocks.environment_v1_service import EnvironmentV1ServiceMock
 from frogml_services_mock.mocks.execution_management_service import (
     ExecutionManagementServiceMock,
 )
@@ -189,6 +203,9 @@ from frogml_services_mock.mocks.location_discovery_service_api import (
     LocationDiscoveryServiceApiMock,
 )
 from frogml_services_mock.mocks.logging_service import LoggingServiceApiMock
+from frogml_services_mock.mocks.model_deployment_manager_service_mock import (
+    ModelDeploymentManagerMock,
+)
 from frogml_services_mock.mocks.model_group_management_service import (
     ModelGroupManagementServiceMock,
 )
@@ -248,13 +265,21 @@ def frogml_container():
         kube_deployment_captain,
         location_discovery,
         logging_client,
+        model_deployment_manager,
         model_group_management,
         model_management,
         model_version_manager,
         secret_service,
         user_application_instance,
     )
-    from frogml.core.clients.administration import eco_system, self_service, authentication
+    from frogml.core.clients.administration import (
+        authentication,
+        cluster_v2,
+        eco_system,
+        environment,
+        environment_v1,
+        self_service,
+    )
     from frogml.core.clients.integration_management import integration_manager_client
     from frogml.core.clients.system_secret import system_secret_client
 
@@ -268,8 +293,11 @@ def frogml_container():
             analytics,
             batch_job_management,
             build_orchestrator,
+            cluster_v2,
             data_versioning,
             deployment,
+            environment,
+            environment_v1,
             instance_template,
             feature_store,
             file_versioning,
@@ -287,6 +315,7 @@ def frogml_container():
             model_version_manager,
             jfrog_gateway,
             location_discovery,
+            model_deployment_manager,
         ]
     )
 
@@ -501,6 +530,26 @@ def attach_servicers(free_port, server):
                 JFrogTenantInfoServiceMock,
                 add_JFrogTenantInfoServiceServicer_to_server,
             ),
+            (
+                "model_deployment_manager_mock",
+                ModelDeploymentManagerMock,
+                add_ModelDeploymentServiceServicer_to_server,
+            ),
+            (
+                "cluster_v2_service_mock",
+                ClusterV2ServiceMock,
+                add_ClusterServiceServicer_to_server,
+            ),
+            (
+                "environment_v0_service_mock",
+                EnvironmentV0ServiceMock,
+                add_EnvironmentV0ServiceServicer_to_server,
+            ),
+            (
+                "environment_v1_service_mock",
+                EnvironmentV1ServiceMock,
+                add_EnvironmentV1ServiceServicer_to_server,
+            ),
             ("port", free_port, None),
         ],
     )