frogml-core 0.0.1__py3-none-any.whl

frogml_core-0.0.1.dist-info/WHEEL

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

frogml_services_mock/__init__.py

@@ -0,0 +1 @@
+from .services_mock import frogml_services_mock

frogml_services_mock/mocks/alert_manager_service_api.py

@@ -0,0 +1,56 @@
+import uuid
+
+from _frogml_proto.qwak.deployment.alert_service_pb2 import (
+    ApplyNotificationChannelRequest,
+    ApplyNotificationChannelResponse,
+    DeleteNotificationChannelRequest,
+    DeleteNotificationChannelResponse,
+    GetNotificationChannelDetailsListRequest,
+    GetNotificationChannelDetailsListResponse,
+)
+from _frogml_proto.qwak.deployment.alert_service_pb2_grpc import (
+    AlertManagementServiceServicer,
+)
+from frogml_services_mock.mocks.utils.exception_handlers import (
+    raise_internal_grpc_error,
+)
+
+
+class AlertManagerServiceApiMock(AlertManagementServiceServicer):
+    def __init__(self):
+        super(AlertManagerServiceApiMock, self).__init__()
+        self.notification_channels = dict()
+
+    def ApplyNotificationChannel(
+        self, request: ApplyNotificationChannelRequest, context
+    ) -> ApplyNotificationChannelResponse:
+        """Channels to add or update"""
+        try:
+            request.notification_channel.notification_channel_id = str(uuid.uuid4())
+            self.notification_channels[
+                request.notification_channel.notification_channel_id
+            ] = request.notification_channel
+            return ApplyNotificationChannelResponse()
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def DeleteNotificationChannel(
+        self, request: DeleteNotificationChannelRequest, context
+    ) -> DeleteNotificationChannelResponse:
+        """Delete notification channel"""
+        try:
+            del self.notification_channels[request.notification_channel_id]
+            return DeleteNotificationChannelResponse()
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def GetNotificationChannelDetailsList(
+        self, request: GetNotificationChannelDetailsListRequest, context
+    ) -> GetNotificationChannelDetailsListResponse:
+        """Get Full Details for channels"""
+        try:
+            return GetNotificationChannelDetailsListResponse(
+                notification_channels=list(self.notification_channels.values())
+            )
+        except Exception as e:
+            raise_internal_grpc_error(context, e)

frogml_services_mock/mocks/alert_registry_service_api.py

@@ -0,0 +1,62 @@
+import uuid
+
+from _frogml_proto.qwak.monitoring.v0.alerting_channel_management_service_pb2 import (
+    CreateAlertingChannelRequest,
+    CreateAlertingChannelResponse,
+    DeleteAlertingChannelRequest,
+    DeleteAlertingChannelResponse,
+    GetAlertingChannelRequestByName,
+    GetAlertingChannelResponseByName,
+    ListAlertingChannelRequest,
+    ListAlertingChannelResponse,
+)
+from _frogml_proto.qwak.monitoring.v0.alerting_channel_management_service_pb2_grpc import (
+    AlertingChannelManagementServiceServicer,
+)
+from _frogml_proto.qwak.monitoring.v0.alerting_channel_pb2 import (
+    AlertingChannelDescription,
+    AlertingChannelMetadata,
+)
+from frogml_services_mock.mocks.utils.exception_handlers import (
+    raise_internal_grpc_error,
+)
+
+
+class AlertsRegistryServiceApiMock(AlertingChannelManagementServiceServicer):
+    def __init__(self):
+        super(AlertsRegistryServiceApiMock, self).__init__()
+        self.alerts_channels = dict()
+
+    def CreateAlertingChannel(self, request: CreateAlertingChannelRequest, context):
+        metadata = AlertingChannelMetadata(
+            id=str(uuid.uuid4()), name=request.options.name
+        )
+        self.alerts_channels[request.options.name] = AlertingChannelDescription(
+            spec=request.options.spec, metadata=metadata
+        )
+        return CreateAlertingChannelResponse()
+
+    def GetAlertingChannelByName(
+        self, request: GetAlertingChannelRequestByName, context
+    ):
+        try:
+            return GetAlertingChannelResponseByName(
+                description=self.alerts_channels[request.name]
+            )
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def ListAlertingChannel(self, request: ListAlertingChannelRequest, context):
+        try:
+            return ListAlertingChannelResponse(
+                description=[desc for name, desc in self.alerts_channels.items()]
+            )
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def DeleteAlertingChannel(self, request: DeleteAlertingChannelRequest, context):
+        try:
+            del self.alerts_channels[request.name]
+            return DeleteAlertingChannelResponse()
+        except Exception as e:
+            raise_internal_grpc_error(context, e)

frogml_services_mock/mocks/analytics_api.py

@@ -0,0 +1,54 @@
+from _frogml_proto.qwak.analytics.analytics_pb2 import QueryStatus
+from _frogml_proto.qwak.analytics.analytics_service_pb2 import (
+    GetQueryResultDownloadURLResponse,
+    GetQueryResultsPageResponse,
+    QueryResponse,
+)
+from _frogml_proto.qwak.analytics.analytics_service_pb2_grpc import (
+    AnalyticsQueryServiceServicer,
+)
+
+SUCCESSFUL_QUERY = "select * from successful"
+FAILED_QUERY = "select * from failed"
+TIMEOUT_QUERY = "select * from timeout"
+CANCELLED_QUERY = "select * from cancelled"
+SUCCESSFUL_QUERY_ID = "successful"
+FAILED_QUERY_ID = "failed"
+TIMEOUT_QUERY_ID = "timeout"
+CANCELLED_QUERY_ID = "cancelled"
+FAILURE_REASON = "failure reason"
+DOWNLOAD_URL = "https://qwak.com/download"
+
+
+class AnalyticsApiMock(AnalyticsQueryServiceServicer):
+    def __init__(self):
+        super().__init__()
+
+    def Query(self, request, context):
+        query = request.query
+        if query == SUCCESSFUL_QUERY:
+            return QueryResponse(query_id=SUCCESSFUL_QUERY_ID)
+        elif query == FAILED_QUERY:
+            return QueryResponse(query_id=FAILED_QUERY_ID)
+        elif query == TIMEOUT_QUERY:
+            return QueryResponse(query_id=TIMEOUT_QUERY_ID)
+        elif query == CANCELLED_QUERY:
+            return QueryResponse(query_id=CANCELLED_QUERY_ID)
+        else:
+            return QueryResponse(query_id="unknown")
+
+    def GetQueryResultsPage(self, request, context):
+        query_id = request.query_id
+        if query_id == SUCCESSFUL_QUERY_ID:
+            return GetQueryResultsPageResponse(status=QueryStatus.SUCCESS, data=None)
+        elif query_id == FAILED_QUERY_ID:
+            return GetQueryResultsPageResponse(
+                status=QueryStatus.FAILED, failure_reason=FAILURE_REASON
+            )
+        elif query_id == TIMEOUT_QUERY_ID:
+            return GetQueryResultsPageResponse(status=QueryStatus.PENDING, data=None)
+        elif query_id == CANCELLED_QUERY_ID:
+            return GetQueryResultsPageResponse(status=QueryStatus.CANCELED, data=None)
+
+    def GetQueryResultDownloadURL(self, request, context):
+        return GetQueryResultDownloadURLResponse(download_url=DOWNLOAD_URL)

frogml_services_mock/mocks/audience_service_api.py

@@ -0,0 +1,79 @@
+import uuid
+
+from _frogml_proto.qwak.audience.v1.audience_api_pb2 import (
+    CreateAudienceRequest,
+    CreateAudienceResponse,
+    DeleteAudienceRequest,
+    DeleteAudienceResponse,
+    GetAudienceRequest,
+    GetAudienceResponse,
+    ListAudienceRequest,
+    ListAudienceResponse,
+    SyncAudiencesRequest,
+    SyncAudiencesResponse,
+    UpdateAudienceRequest,
+    UpdateAudienceResponse,
+)
+from _frogml_proto.qwak.audience.v1.audience_api_pb2_grpc import AudienceAPIServicer
+from _frogml_proto.qwak.audience.v1.audience_pb2 import AudienceEntry
+from frogml_services_mock.mocks.utils.exception_handlers import (
+    raise_internal_grpc_error,
+)
+
+
+class AudienceServiceApiMock(AudienceAPIServicer):
+    def __init__(self):
+        super(AudienceServiceApiMock, self).__init__()
+        self.audiences = dict()
+
+    def CreateAudience(
+        self, request: CreateAudienceRequest, context
+    ) -> CreateAudienceResponse:
+        try:
+            audience_id = str(uuid.uuid4())
+            self.audiences[audience_id] = request.audience
+            return CreateAudienceResponse(audience_id=audience_id)
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def UpdateAudience(
+        self, request: UpdateAudienceRequest, context
+    ) -> UpdateAudienceResponse:
+        try:
+            self.audiences[request.id] = request.audience
+            return UpdateAudienceResponse()
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def GetAudience(self, request: GetAudienceRequest, context) -> GetAudienceResponse:
+        try:
+            return GetAudienceResponse(audience=self.audiences[request.audience_id])
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def ListAudience(
+        self, request: ListAudienceRequest, context
+    ) -> ListAudienceResponse:
+        try:
+            return ListAudienceResponse(
+                audience_entries=[
+                    AudienceEntry(id=audience_id, audience=audience)
+                    for audience_id, audience in self.audiences.items()
+                ]
+            )
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def DeleteAudience(
+        self, request: DeleteAudienceRequest, context
+    ) -> DeleteAudienceResponse:
+        try:
+            del self.audiences[request.audience_id]
+            return DeleteAudienceResponse()
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def SyncAudiences(
+        self, request: SyncAudiencesRequest, context
+    ) -> SyncAudiencesResponse:
+        return SyncAudiencesResponse()

frogml_services_mock/mocks/authentication_service.py

@@ -0,0 +1,33 @@
+from datetime import datetime, timedelta
+
+from jose import jwt
+
+from _frogml_proto.qwak.administration.v0.authentication.authentication_service_pb2 import (
+    AuthenticateRequest,
+    AuthenticateResponse,
+)
+from _frogml_proto.qwak.administration.v0.authentication.authentication_service_pb2_grpc import (
+    AuthenticationServiceServicer,
+)
+
+epoch = datetime.timestamp(datetime.utcnow() + timedelta(days=1))
+
+_payload = {
+    "https://auth-token.qwak.ai/qwak-partner-id": "9",
+    "https://auth-token.qwak.ai/qwak-user-id": "Test User",
+    "https://auth-token.qwak.ai/mlflow-server": "http://localhost",
+    "https://auth-token.qwak.ai/data-api": "http://localhost",
+    "https://auth-token.qwak.ai/grpc-api": "localhost",
+    "https://auth-token.qwak.ai/models-api": "localhost",
+    "https://auth-token.qwak.ai/qwak-bucket": "frogml.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)

frogml_services_mock/mocks/automation_management_service.py

@@ -0,0 +1,228 @@
+import itertools
+import uuid
+from collections import defaultdict
+from datetime import datetime
+
+from _frogml_proto.qwak.automation.v1.automation_management_service_pb2 import (
+    CreateAutomationRequest,
+    CreateAutomationResponse,
+    DeleteAutomationRequest,
+    DeleteAutomationResponse,
+    GetAutomationByNameRequest,
+    GetAutomationByNameResponse,
+    GetAutomationRequest,
+    GetAutomationResponse,
+    ListAutomationExecutionsRequest,
+    ListAutomationExecutionsResponse,
+    ListAutomationsRequest,
+    ListAutomationsResponse,
+    RegisterAutomationExecutionRequest,
+    RegisterAutomationExecutionResponse,
+    RunAutomationRequest,
+    RunAutomationResponse,
+    ToggleAutomationActivityRequest,
+    ToggleAutomationActivityResponse,
+    UpdateAutomationExecutionRequest,
+    UpdateAutomationExecutionResponse,
+    UpdateAutomationRequest,
+    UpdateAutomationResponse,
+)
+from _frogml_proto.qwak.automation.v1.automation_management_service_pb2_grpc import (
+    AutomationManagementServiceServicer,
+)
+from _frogml_proto.qwak.automation.v1.automation_pb2 import Automation
+from frogml.automations.automation_executions import (
+    AutomationExecution,
+    ExecutionRunDetails,
+    ExecutionStatus,
+)
+from frogml_services_mock.mocks.utils.exception_handlers import (
+    raise_internal_grpc_error,
+)
+
+
+class AutomationManagementServiceMock(AutomationManagementServiceServicer):
+    FROGML_ENVIRONMENT = "test_environment"
+
+    def __init__(self):
+        self.automations = dict()
+        self.automation_executions = defaultdict(list)
+
+    def CreateAutomation(
+        self, request: CreateAutomationRequest, context
+    ) -> CreateAutomationResponse:
+        """Create an automation"""
+        try:
+            automation_id = str(uuid.uuid4())
+            automation = Automation(
+                automation_id=automation_id,
+                automation_spec=request.automation_spec,
+                qwak_environment_id=self.FROGML_ENVIRONMENT,
+            )
+            self.automations[automation_id] = automation
+            return CreateAutomationResponse(automation_id=automation_id)
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def UpdateAutomation(
+        self, request: UpdateAutomationRequest, context
+    ) -> UpdateAutomationResponse:
+        """Update an automation"""
+        try:
+            automation = Automation(
+                automation_id=request.automation_id,
+                automation_spec=request.automation_spec,
+                qwak_environment_id=self.FROGML_ENVIRONMENT,
+            )
+            self.automations[request.automation_id] = automation
+            return UpdateAutomationResponse()
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def DeleteAutomation(
+        self, request: DeleteAutomationRequest, context
+    ) -> DeleteAutomationResponse:
+        """Delete an automation"""
+        try:
+            del self.automations[request.automation_id]
+            return DeleteAutomationResponse()
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def GetAutomation(
+        self, request: GetAutomationRequest, context
+    ) -> GetAutomationResponse:
+        """Get an automation By ID"""
+        try:
+            return GetAutomationResponse(
+                automation=self.automations[request.automation_id]
+            )
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def GetAutomationByName(
+        self, request: GetAutomationByNameRequest, context
+    ) -> GetAutomationByNameResponse:
+        """Get an automation by name"""
+        try:
+            return GetAutomationByNameResponse(
+                automation=list(
+                    filter(
+                        lambda a: a.automation_spec.automation_name
+                        == request.automation_name,
+                        self.automations.values(),
+                    )
+                )[0]
+            )
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def ListAutomations(
+        self, request: ListAutomationsRequest, context
+    ) -> ListAutomationsResponse:
+        """List all automation"""
+        try:
+            return ListAutomationsResponse(automations=list(self.automations.values()))
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def ListAutomationExecutions(
+        self, request: ListAutomationExecutionsRequest, context
+    ) -> ListAutomationExecutionsResponse:
+        """List automation's executions"""
+        try:
+            return ListAutomationExecutionsResponse(
+                automation_executions=self.automation_executions[request.automation_id]
+            )
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def RegisterAutomationExecution(
+        self, request: RegisterAutomationExecutionRequest, context
+    ) -> RegisterAutomationExecutionResponse:
+        """Register an execution"""
+        try:
+            run_details = ExecutionRunDetails(
+                start_time=datetime.now(), status=ExecutionStatus.RUNNING
+            )
+            execution = AutomationExecution(
+                execution_id=request.execution_id, run_details=run_details
+            )
+            self.automation_executions[request.automation_id].append(execution)
+            return RegisterAutomationExecutionResponse()
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def UpdateAutomationExecution(
+        self, request: UpdateAutomationExecutionRequest, context
+    ) -> UpdateAutomationExecutionResponse:
+        """Update an execution"""
+        try:
+            existing_executions = list(
+                itertools.chain(*self.automation_executions.values())
+            )
+            relevant_execution = list(
+                filter(
+                    lambda ex: ex.execution_id == request.execution_id,
+                    existing_executions,
+                )
+            )[0]
+            self.__merge_run_details(
+                relevant_execution.run_details, request.run_details
+            )
+            return UpdateAutomationExecutionResponse()
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def ToggleAutomationActivity(
+        self, request: ToggleAutomationActivityRequest, context
+    ) -> ToggleAutomationActivityResponse:
+        """Toggle whether an automation is enabled or disabled"""
+        try:
+            self.automations[request.automation_id].automation_spec = request.is_enabled
+            return ToggleAutomationActivityResponse()
+        except Exception as e:
+            raise_internal_grpc_error(context, e)
+
+    def RunAutomation(
+        self, request: RunAutomationRequest, context
+    ) -> RunAutomationResponse:
+        """Run an automation"""
+        if self.automations[request.automation_id]:
+            return RunAutomationResponse()
+        else:
+            raise_internal_grpc_error(
+                context,
+                ValueError(f"Automation with ID {request.automation_id} not found"),
+            )
+
+    @staticmethod
+    def __merge_run_details(
+        run_details_1: ExecutionRunDetails, run_details_2: ExecutionRunDetails
+    ):
+        run_details_1.start_time = (
+            run_details_2.start_time
+            if run_details_2.start_time
+            else run_details_1.start_time
+        )
+        run_details_1.end_time = (
+            run_details_2.end_time if run_details_2.end_time else run_details_1.end_time
+        )
+        run_details_1.error_details = (
+            run_details_2.error_details
+            if run_details_2.error_details
+            else run_details_1.error_details
+        )
+        run_details_1.status = (
+            run_details_2.status if run_details_2.status else run_details_1.status
+        )
+        run_details_1.task = (
+            run_details_2.task if run_details_2.task else run_details_1.task
+        )
+        run_details_1.finish_cause = (
+            run_details_2.finish_cause
+            if run_details_2.finish_cause
+            else run_details_1.finish_cause
+        )
+
+        return run_details_1

frogml_services_mock/mocks/autoscaling_service_api.py

@@ -0,0 +1,31 @@
+import uuid
+
+from _frogml_proto.qwak.auto_scaling.v1.auto_scaling_service_pb2 import (
+    AttachAutoScalingRequest,
+    AttachAutoScalingResponse,
+)
+from _frogml_proto.qwak.auto_scaling.v1.auto_scaling_service_pb2_grpc import (
+    AutoScalingServiceServicer,
+)
+from frogml_services_mock.mocks.utils.exception_handlers import (
+    raise_internal_grpc_error,
+)
+
+
+class AutoscalingServiceApiMock(AutoScalingServiceServicer):
+    def __init__(self):
+        super(AutoscalingServiceApiMock, self).__init__()
+        self.autoscaling_policies = dict()
+
+    def AttachAutoScaling(
+        self, request: AttachAutoScalingRequest, context
+    ) -> AttachAutoScalingResponse:
+        try:
+            autoscale_id = str(uuid.uuid4())
+            self.autoscaling_policies[request.model_id] = (
+                autoscale_id,
+                request.auto_scaling_config,
+            )
+            return AttachAutoScalingResponse(auto_scaling_id=autoscale_id)
+        except Exception as e:
+            raise_internal_grpc_error(context, e)