qwak-core 0.1.37__py3-none-any.whl → 0.5.19__py3-none-any.whl

{qwak_core-0.1.37.dist-info → qwak_core-0.5.19.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 1.4.0
+Generator: poetry-core 2.1.3
 Root-Is-Purelib: true
 Tag: py3-none-any

qwak_services_mock/mocks/authentication_service.py

@@ -1,3 +1,5 @@
+from datetime import datetime, timedelta
+
 from _qwak_proto.qwak.administration.v0.authentication.authentication_service_pb2 import (
     AuthenticateRequest,
     AuthenticateResponse,
@@ -7,6 +9,7 @@ from _qwak_proto.qwak.administration.v0.authentication.authentication_service_pb
 )
 from jose import jwt
 
+epoch = datetime.timestamp(datetime.utcnow() + timedelta(days=1))
 
 _payload = {
     "https://auth-token.qwak.ai/qwak-partner-id": "9",
@@ -16,16 +19,14 @@ _payload = {
     "https://auth-token.qwak.ai/grpc-api": "localhost",
     "https://auth-token.qwak.ai/models-api": "localhost",
     "https://auth-token.qwak.ai/qwak-bucket": "qwak.bucket",
+    "exp": epoch,
 }
 
 mock_jwt_token = jwt.encode(_payload, "secret", algorithm="HS256")
 
 
 class AuthenticationServiceMock(AuthenticationServiceServicer):
-
     def Authenticate(
         self, request: AuthenticateRequest, context
     ) -> AuthenticateResponse:
-        return AuthenticateResponse(
-            access_token=mock_jwt_token
-        )
+        return AuthenticateResponse(access_token=mock_jwt_token)

qwak_services_mock/mocks/batch_job_manager_service.py

@@ -9,6 +9,7 @@ from _qwak_proto.qwak.batch_job.v1.batch_job_service_pb2 import (
     BatchJobDetails,
     BatchJobMessage,
     BatchJobStatusMessage,
+    BatchTaskStatusMessage,
     CancelBatchJobRequest,
     CancelBatchJobResponse,
     CancelWarmupJobRequest,
@@ -34,6 +35,10 @@ from _qwak_proto.qwak.batch_job.v1.batch_job_service_pb2 import (
     StartBatchJobResponse,
     StartWarmupJobRequest,
     StartWarmupJobResponse,
+    TaskExecutionDetails,
+    UpdateTasksDetailsRequest,
+    UpdateTasksDetailsResponse,
+    BatchTaskDetails,
 )
 from _qwak_proto.qwak.batch_job.v1.batch_job_service_pb2_grpc import (
     BatchJobManagementServiceServicer,
@@ -49,6 +54,7 @@ class BatchJobManagerService(BatchJobManagementServiceServicer):
         self.id_to_batch_job: Dict[str, MockBatchJob] = dict()
         self.model_id_active_warmup: Set[str] = set()
         self.models_to_fail: Set[str] = set()
+        self.task_id_to_update_task_details: Dict[str, BatchTaskDetails] = dict()
 
     def StartBatchJob(
         self, request: StartBatchJobRequest, context
@@ -144,12 +150,21 @@ class BatchJobManagerService(BatchJobManagementServiceServicer):
         self, request: ListBatchJobsRequest, context
     ) -> ListBatchJobsResponse:
         try:
-            batch_jobs = list(
-                filter(
-                    lambda job: job.start_request.model_id == request.model_id,
-                    self.id_to_batch_job.values(),
+            if request.build_id:
+                batch_jobs = list(
+                    filter(
+                        lambda job: job.start_request.build_id == request.build_id,
+                        self.id_to_batch_job.values(),
+                    )
                 )
-            )
+            else:
+                batch_jobs = list(
+                    filter(
+                        lambda job: job.start_request.model_id == request.model_id,
+                        self.id_to_batch_job.values(),
+                    )
+                )
+
             batch_jobs_details = list(map(batch_job_to_details, batch_jobs))
             return ListBatchJobsResponse(batch_jobs=batch_jobs_details)
         except Exception as e:
@@ -254,6 +269,25 @@ class BatchJobManagerService(BatchJobManagementServiceServicer):
         except Exception as e:
             raise_internal_grpc_error(context, e)
 
+    def UpdateTasksDetails(
+            self, request: UpdateTasksDetailsRequest, context
+    ) -> UpdateTasksDetailsResponse:
+        for task in request.tasks_details:
+            task_id = task.task_id
+            if not task_id :
+                raise ValueError("Task ID cannot be empty")
+            input_files = task.input_files_details
+
+            print(f"Updating task with ID: {task_id} and input files: {[file.path for file in input_files]}")
+            self.task_id_to_update_task_details[task_id] = task
+        return UpdateTasksDetailsResponse()
+
+    def get_task_details(self,task_id: str
+        ) -> Optional[BatchTaskDetails]:
+            """
+            Get task details by task ID
+            """
+            return self.task_id_to_update_task_details[task_id]
 
 class MockBatchJob:
     def __init__(self, request: StartBatchJobRequest):
@@ -267,10 +301,22 @@ class MockBatchJob:
         self.finished_files: int = 0
         self.total_files: int = 0
         self.report: List[str] = []
+        self.tasks: List[TaskExecutionDetails] = []
 
     def add_report_messages(self, messages: List[str]):
         self.report.extend(messages)
 
+    def add_task_executions(self):
+        self.tasks.append(
+            TaskExecutionDetails(
+                task_id=str(uuid.uuid4()),
+                status=BatchTaskStatusMessage.BATCH_TASK_RUNNING_STATUS,
+                start_time=Timestamp(),
+                end_time=Timestamp(),
+                filename="test",
+            )
+        )
+
 
 def batch_job_to_details(batch_job: MockBatchJob) -> BatchJobDetails:
     return BatchJobDetails(
@@ -307,4 +353,5 @@ def batch_job_to_message(batch_job: MockBatchJob) -> BatchJobMessage:
             ExecutionReportLine(time=Timestamp(), line=line)
             for line in batch_job.report
         ],
-    )
+        task_executions=batch_job.tasks,
+    )

qwak_services_mock/mocks/build_orchestrator_build_api.py

@@ -1,3 +1,5 @@
+from typing import List
+
 from _qwak_proto.qwak.build.v1.build_api_pb2 import (
     GetBuildRequest,
     GetBuildResponse,
@@ -13,12 +15,18 @@ from _qwak_proto.qwak.build.v1.build_api_pb2 import (
     RegisterModelSchemaResponse,
     RegisterTagsRequest,
     RegisterTagsResponse,
+    SaveFrameworkModelsRequest,
+    SaveFrameworkModelsResponse,
     UpdateBuildStatusRequest,
     UpdateBuildStatusResponse,
 )
 from _qwak_proto.qwak.build.v1.build_api_pb2_grpc import BuildAPIServicer
 from _qwak_proto.qwak.build.v1.build_pb2 import Audit, Build, BuildStatus
 from google.protobuf.timestamp_pb2 import Timestamp
+from qwak.clients.build_orchestrator.client import (
+    FrameworkModelDataClass,
+    HuggingModelDataClass,
+)
 from qwak_services_mock.mocks.utils.exception_handlers import raise_internal_grpc_error
 
 
@@ -26,10 +34,15 @@ class BuildOrchestratorBuildApiMock(BuildAPIServicer):
     def __init__(self):
         super(BuildOrchestratorBuildApiMock, self).__init__()
         self._builds: dict[str:Build] = {}
+        self._framework_models: dict[str : List[FrameworkModelDataClass]] = {}
 
     def given_build(self, build: Build):
         self._builds[build.buildId] = build
 
+    def SaveFrameworkModels(self, request: SaveFrameworkModelsRequest, context):
+        self._framework_models[request.spec.build_id] = request.spec.framework_models
+        return SaveFrameworkModelsResponse()
+
     def RegisterBuild(
         self, request: RegisterBuildRequest, context
     ) -> RegisterBuildResponse:
@@ -82,7 +95,11 @@ class BuildOrchestratorBuildApiMock(BuildAPIServicer):
         return GetBuildResponse(build=self._builds[request.build_id])
 
     def ListBuilds(self, request: ListBuildsRequest, context) -> ListBuildsResponse:
-        valid_builds = [self._builds[buildId] for buildId in self._builds if self._builds[buildId].model_uuid == request.model_uuid]
+        valid_builds = [
+            self._builds[buildId]
+            for buildId in self._builds
+            if self._builds[buildId].model_uuid == request.model_uuid
+        ]
         builds_to_return = []
 
         if len(request.filter.tags) > 0:

qwak_services_mock/mocks/build_orchestrator_build_settings_api.py

@@ -0,0 +1,35 @@
+from _qwak_proto.qwak.build_settings.build_settings_api_pb2 import (
+    DeleteBuildSettingsRequest,
+    DeleteBuildSettingsResponse,
+    GetBuildSettingsRequest,
+    GetBuildSettingsResponse,
+    SaveBuildSettingsRequest,
+    SaveBuildSettingsResponse,
+)
+from _qwak_proto.qwak.build_settings.build_settings_api_pb2_grpc import (
+    BuildSettingsApiServicer,
+)
+from _qwak_proto.qwak.build_settings.build_settings_pb2 import BuildSettings
+from google.protobuf.timestamp_pb2 import Timestamp
+from qwak_services_mock.mocks.utils.exception_handlers import raise_internal_grpc_error
+
+
+class BuildOrchestratorBuildSettingsApiMock(BuildSettingsApiServicer):
+    def __init__(self):
+        self._build_settings = dict()
+
+    def given_build_settings(self, environment_id: str, build_settings: BuildSettings):
+        self._build_settings[environment_id] = build_settings
+
+    def GetBuildSettings(
+        self, request: GetBuildSettingsRequest, context
+    ) -> GetBuildSettingsResponse:
+        build_settings = self._build_settings.get(request.environment_id)
+        if not build_settings:
+            raise_internal_grpc_error(
+                context,
+                Exception(
+                    f"Build settings for environment id {request.environment_id} doesn't exist"
+                ),
+            )
+        return GetBuildSettingsResponse(build_settings=build_settings)

qwak_services_mock/mocks/build_orchestrator_service_api.py

@@ -65,6 +65,10 @@ class BuildOrchestratorServiceApiMock(BuildsOrchestratorServiceServicer):
     def GetBuildVersioningDownloadURL(self, request, context):
         """Request the service for a path ready for download the artifact"""
         try:
+            if request.params.build_id == "not_found":
+                context.set_code(grpc.StatusCode.NOT_FOUND)
+                context.set_details(str("key not found"))
+                return GetBuildVersioningDownloadURLResponse()
             return GetBuildVersioningDownloadURLResponse(
                 download_url=BuildOrchestratorServiceApiMock.DOWNLOAD_URL_FORMAT.format(
                     build_id=request.params.build_id

qwak_services_mock/mocks/data_versioning_service.py

@@ -23,7 +23,7 @@ class DataVersioningServiceMock(DataVersioningManagementServiceServicer):
         self, request: RegisterDataTagRequest, context
     ) -> RegisterDataTagResponse:
         try:
-            self.tags[request.data_tag_spec.model_id].append(request.data_tag_spec)
+            self.tags[request.data_tag_spec.build_id].append(request.data_tag_spec)
             return RegisterDataTagResponse()
         except Exception as e:
             raise_internal_grpc_error(context, e)
@@ -32,6 +32,30 @@ class DataVersioningServiceMock(DataVersioningManagementServiceServicer):
         self, request: GetModelDataTagsRequest, context
     ) -> GetModelDataTagsResponse:
         try:
-            return GetModelDataTagsResponse(data_tags=self.tags[request.model_id])
+            data_tags_list = []
+
+            if not request.build_id:
+                for all_data_tags in self.tags.values():
+                    for data_tag in all_data_tags:
+                        data_tags_list.append(data_tag)
+            else:
+                for data_tag_by_build_id in self.tags[request.build_id]:
+                    if data_tag_by_build_id.model_id == request.model_id:
+                        data_tags_list.append(data_tag_by_build_id)
+
+            if request.filter:
+                filter_type: str = request.filter.WhichOneof("filter")
+
+                if filter_type == "tag_contains":
+                    for data_tag in data_tags_list:
+                        if request.filter.tag_contains not in data_tag.tag:
+                            data_tags_list.remove(data_tag)
+
+                if filter_type == "tag_prefix":
+                    for data_tag in data_tags_list:
+                        if not data_tag.tag.startswith(request.filter.tag_prefix):
+                            data_tags_list.remove(data_tag)
+
+            return GetModelDataTagsResponse(data_tags=data_tags_list)
         except Exception as e:
             raise_internal_grpc_error(context, e)

qwak_services_mock/mocks/deployment_management_service.py

@@ -3,6 +3,7 @@ from collections import defaultdict
 from typing import Dict, Optional
 
 from _qwak_proto.qwak.deployment.deployment_pb2 import (
+    DeploymentBrief,
     DeploymentDetails,
     EnvironmentApplyModelTrafficResponse,
     EnvironmentDeploymentDetailsMessage,
@@ -23,6 +24,8 @@ from _qwak_proto.qwak.deployment.deployment_service_pb2 import (
     DeployModelResponse,
     GetDeploymentDetailsRequest,
     GetDeploymentDetailsResponse,
+    GetDeploymentHistoryRequest,
+    GetDeploymentHistoryResponse,
     GetDeploymentStatusRequest,
     GetDeploymentStatusResponse,
     GetModelTrafficRequest,
@@ -117,6 +120,9 @@ class DeploymentManagementServiceMock(DeploymentManagementServiceServicer):
                     self.model_uuid_to_active_deployment_ids[
                         env_request_details.model_uuid
                     ].append(deployment_id)
+                    self.deployments_history[env_request_details.model_uuid].insert(
+                        0, self.active_deployments[deployment_key]
+                    )
                     if (
                         env_request_details.hosting_service.kube_deployment.kube_deployment_type
                         == KubeDeploymentType.ONLINE
@@ -232,6 +238,20 @@ class DeploymentManagementServiceMock(DeploymentManagementServiceServicer):
         except Exception as e:
             raise_internal_grpc_error(context, e)
 
+    def GetDeploymentHistory(
+        self, request: GetDeploymentHistoryRequest, context
+    ) -> GetDeploymentHistoryResponse:
+        try:
+            result = []
+            for deployment_details in self.deployments_history[request.model_uuid]:
+                deployment_brief_proto: DeploymentBrief = (
+                    self.__deployment_brief_to_proto(deployment_details)
+                )
+                result.append(deployment_brief_proto)
+            return GetDeploymentHistoryResponse(deployments=result)
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
     def UpdateDeploymentRuntimeSettings(
         self, request: UpdateDeploymentRuntimeSettingsRequest, context
     ) -> UpdateDeploymentRuntimeSettingsResponse:
@@ -406,25 +426,33 @@ class DeploymentManagementServiceMock(DeploymentManagementServiceServicer):
             deployment_id=deployment_details["deployment_id"],
         )
 
-    def given_deployment(self, model_id: str,
-                environment_id: str, status: int, model_uuid: str, deployment_id: str) -> str:
-
-        deployment_details = (
-            self.active_deployments.get(
-                (model_id, "", environment_id)
-            )
+    def __deployment_brief_to_proto(self, deployment_details):
+        return DeploymentBrief(
+            build_id=deployment_details["build_id"],
+            status=deployment_details["status"],
+            deployment_timestamp=deployment_details["timestamp"],
+            environment_id=deployment_details["env"],
+            variation_name=deployment_details[
+                "hosting_service"
+            ].kube_deployment.serving_strategy.realtime_config.traffic_config.selected_variation_name,
+            variation_protect_state=deployment_details[
+                "hosting_service"
+            ].kube_deployment.serving_strategy.realtime_config.traffic_config.selected_variation_protect_state,
+            deployment_id=deployment_details["deployment_id"],
         )
+
+    def given_deployment(
+        self,
+        model_id: str,
+        environment_id: str,
+        status: int,
+        model_uuid: str,
+        deployment_id: str,
+    ) -> str:
+        deployment_details = self.active_deployments.get((model_id, "", environment_id))
         deployment_details["status"] = status
         deployment_details["model_uuid"] = [model_uuid]
-        self.active_deployments[
-            (model_id, "", environment_id)
-        ] = deployment_details
-        self.model_uuid_to_active_deployment_ids = defaultdict(
-            list
-        )
-        self.model_uuid_to_active_deployment_ids[
-            model_uuid
-        ] = [
-            deployment_id
-        ]
+        self.active_deployments[(model_id, "", environment_id)] = deployment_details
+        self.model_uuid_to_active_deployment_ids = defaultdict(list)
+        self.model_uuid_to_active_deployment_ids[model_uuid] = [deployment_id]
         return deployment_id

qwak_services_mock/mocks/ecosystem_service_api.py

@@ -14,7 +14,6 @@ from qwak_services_mock.mocks.utils.exception_handlers import raise_internal_grp
 
 class EcoSystemServiceMock(QwakEcosystemRuntimeServicer):
     def __init__(self):
-
         self._aws_temp_credentials = None
         self._authenticated_user_context = None
         super(EcoSystemServiceMock, self).__init__()

qwak_services_mock/mocks/execution_management_service.py

@@ -1,7 +1,8 @@
-from _qwak_proto.qwak.execution.v1.execution_service_pb2 import TriggerBackfillResponse
+from _qwak_proto.qwak.execution.v1.execution_service_pb2 import (TriggerBackfillResponse, TriggerStreamingAggregationBackfillResponse)
 from _qwak_proto.qwak.execution.v1.execution_service_pb2_grpc import (
     FeatureStoreExecutionServiceServicer,
 )
+
 from grpc import RpcError
 
 
@@ -23,3 +24,10 @@ class ExecutionManagementServiceMock(FeatureStoreExecutionServiceServicer):
         if self._raise_exception_on_request:
             raise RpcError
         return TriggerBackfillResponse(execution_id=self._execution_id)
+
+    def TriggerStreamingAggregationBackfill(self, request, context):
+        if self._raise_exception_on_request:
+            raise RpcError
+        return TriggerStreamingAggregationBackfillResponse(
+            execution_id=self._execution_id
+        )

qwak_services_mock/mocks/feature_store_data_sources_manager_api.py

@@ -1,31 +1,64 @@
+import time
+import uuid
+from typing import Dict, Optional
+
 import grpc
 from _qwak_proto.qwak.feature_store.sources.data_source_pb2 import (
     DataSource,
     DataSourceDefinition,
+    DataSourceMetadata,
+    DataSourceSpec,
+    FeatureSetBrief,
 )
 from _qwak_proto.qwak.feature_store.sources.data_source_service_pb2 import (
     CreateDataSourceResponse,
+    CreateDataSourceUploadURLRequest,
+    CreateDataSourceUploadURLResponse,
+    DeleteDataSourceResponse,
+    GetDataSourceByIdRequest,
+    GetDataSourceByIdResponse,
     GetDataSourceByNameResponse,
+    ListDataSourcesResponse,
 )
 from _qwak_proto.qwak.feature_store.sources.data_source_service_pb2_grpc import (
     DataSourceServiceServicer,
 )
+from _qwak_proto.qwak.feature_store.v1.internal.data_source.data_source_service_pb2 import (
+    GetDataSourceSourceCodeUploadResponse,
+)
 
 
 class DataSourceServiceMock(DataSourceServiceServicer):
     def __init__(self):
-        self._data_sources_spec = {}
+        self._data_sources_spec: Dict[str, DataSourceSpec] = {}
+        self._data_source_id_to_name: Dict[str, str] = {}
+        self._ds_definition_by_name: Dict[str, DataSourceDefinition] = {}
 
     def CreateDataSource(self, request, context):
         data_source_type = request.data_source_spec.WhichOneof("type")
+        data_source_id = str(uuid.uuid4())
+
         ds_name = (
             request.data_source_spec.batch_source.name
             if data_source_type == "batch_source"
             else request.data_source_spec.stream_source.name
         )
+        data_source_definition = DataSourceDefinition(
+            data_source_id=data_source_id,
+            data_source_spec=request.data_source_spec,
+        )
 
         self._data_sources_spec[ds_name] = request.data_source_spec
-        return CreateDataSourceResponse()
+        self._data_source_id_to_name[data_source_id] = ds_name
+        self._ds_definition_by_name[ds_name] = data_source_definition
+
+        return CreateDataSourceResponse(
+            data_source=DataSource(
+                data_source_definition=data_source_definition,
+                metadata=None,
+                feature_sets=[],
+            )
+        )
 
     def GetDataSourceByName(self, request, context):
         if request.data_source_name not in self._data_sources_spec:
@@ -42,3 +75,66 @@ class DataSourceServiceMock(DataSourceServiceServicer):
                 feature_sets=[],
             )
         )
+
+    def GetDataSourceById(self, request: GetDataSourceByIdRequest, context):
+        if request.data_source_id not in self._data_source_id_to_name:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            return GetDataSourceByIdResponse()
+
+        data_source_name = self._data_source_id_to_name[request.data_source_id]
+        data_source_spec = self._data_sources_spec[data_source_name]
+
+        return GetDataSourceByIdResponse(
+            data_source=DataSource(
+                data_source_definition=DataSourceDefinition(
+                    data_source_id=request.data_source_id,
+                    data_source_spec=data_source_spec,
+                ),
+                metadata=None,
+                feature_sets=[],
+            )
+        )
+
+    def CreateDataSourceUploadURLV1(
+        self, request: CreateDataSourceUploadURLRequest, context
+    ):
+        return CreateDataSourceUploadURLResponse(
+            upload_url=f"https://data_source_artifacts.s3.amazonaws.com/{request.data_source_name}_{request.object_name}_{int(time.time() * 1000)}"
+        )
+
+    def GetDataSourceSourceCodeUploadURL(self, request, context):
+        if request.data_source_name not in self._data_sources_spec:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            return GetDataSourceByNameResponse()
+
+        return GetDataSourceSourceCodeUploadResponse(
+            upload_url=f"https://data_source_artifacts.s3.amazonaws.com/{request.data_source_name}_{request.object_name}_{int(time.time() * 1000)}",
+            extra_headers={},
+        )
+
+    def ListDataSources(self, request, context):
+        return ListDataSourcesResponse(
+            data_sources=[
+                DataSource(
+                    data_source_definition=ds_definition,
+                    metadata=DataSourceMetadata(),
+                    feature_sets=[FeatureSetBrief()],
+                )
+                for ds_definition in self._ds_definition_by_name.values()
+            ]
+        )
+
+    def DeleteDataSource(self, request, context):
+        ds_name: Optional[str] = self._data_source_id_to_name.get(
+            request.data_source_id, None
+        )
+        if not ds_name:
+            context.set_details(
+                f"DataSource ID {request.data_source_id} doesn't exist'"
+            )
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+
+        self._data_sources_spec.pop(ds_name)
+        self._data_source_id_to_name.pop(request.data_source_id)
+        self._ds_definition_by_name.pop(ds_name)
+        return DeleteDataSourceResponse()

qwak_services_mock/mocks/feature_store_entities_manager_api.py

@@ -1,3 +1,5 @@
+import uuid
+
 import grpc
 from _qwak_proto.qwak.feature_store.entities.entity_pb2 import (
     Entity,
@@ -7,6 +9,9 @@ from _qwak_proto.qwak.feature_store.entities.entity_pb2 import (
 )
 from _qwak_proto.qwak.feature_store.entities.entity_service_pb2 import (
     CreateEntityResponse,
+    DeleteEntityRequest,
+    DeleteEntityResponse,
+    GetEntityByIdResponse,
     GetEntityByNameResponse,
     ListEntitiesResponse,
 )
@@ -18,23 +23,43 @@ from _qwak_proto.qwak.feature_store.entities.entity_service_pb2_grpc import (
 class EntityServiceMock(EntityServiceServicer):
     def __init__(self):
         self._entities_spec = {}
+        self._entity_id_name = {}
 
     def reset_entities(self):
         self._entities_spec.clear()
 
     def CreateEntity(self, request, context):
+        entity_id = str(uuid.uuid4())
         self._entities_spec[request.entity_spec.name] = request.entity_spec
+        self._entity_id_name[entity_id] = request.entity_spec.name
         return CreateEntityResponse()
 
+    def DeleteEntity(self, request: DeleteEntityRequest, context):
+        entity_id: str = request.entity_id
+
+        if entity_id not in self._entity_id_name.keys():
+            context.set_details(f"Entity ID {request.entity_id} doesn't exist'")
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            return
+
+        entity_name: str = self._entity_id_name[entity_id]
+        del self._entity_id_name[entity_id]
+        del self._entities_spec[entity_name]
+
+        return DeleteEntityResponse()
+
     def GetEntityByName(self, request, context):
         if request.entity_name not in self._entities_spec:
             context.set_code(grpc.StatusCode.NOT_FOUND)
             return GetEntityByNameResponse()
 
+        entity_key: str = list(self._entity_id_name.keys())[
+            list(self._entity_id_name.values()).index(request.entity_name)
+        ]
         return GetEntityByNameResponse(
             entity=Entity(
                 entity_definition=EntityDefinition(
-                    entity_id="123",
+                    entity_id=entity_key,
                     entity_spec=self._entities_spec[request.entity_name],
                 ),
                 metadata=None,
@@ -42,16 +67,34 @@ class EntityServiceMock(EntityServiceServicer):
             )
         )
 
+    def GetEntityById(self, request, context):
+        if request.entity_id not in self._entity_id_name:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            return GetEntityByIdResponse()
+
+        entity_name: str = self._entity_id_name[request.entity_id]
+        return GetEntityByIdResponse(
+            entity=Entity(
+                entity_definition=EntityDefinition(
+                    entity_id=request.entity_id,
+                    entity_spec=self._entities_spec[entity_name],
+                ),
+                metadata=None,
+                feature_sets=[],
+            )
+        )
+
     def ListEntities(self, request, context):
         return ListEntitiesResponse(
             entities=[
                 Entity(
                     entity_definition=EntityDefinition(
-                        entity_spec=entity_spec, entity_id="123"
+                        entity_spec=self._entities_spec[entity_name],
+                        entity_id=entity_id,
                     ),
                     metadata=EntityMetadata(),
                     feature_sets=[FeatureSetBrief()],
                 )
-                for entity_spec in self._entities_spec.values()
+                for entity_id, entity_name in self._entity_id_name.items()
             ]
         )