codemie-test-harness 0.1.222__py3-none-any.whl → 0.1.224__py3-none-any.whl

codemie_test_harness/tests/conftest.py

@@ -951,11 +951,21 @@ def integration(integration_utils):
 
 
 @pytest.fixture(scope="session")
-def vendor_utils():
+def vendor_assistant_utils():
     """Create VendorUtils instance for managing vendor assistants"""
-    from codemie_test_harness.tests.utils.vendor_utils import VendorUtils
+    from codemie_test_harness.tests.utils.vendor_utils import VendorAssistantUtils
 
-    return VendorUtils()
+    return VendorAssistantUtils()
+
+
+@pytest.fixture(scope="session")
+def vendor_workflow_utils():
+    """Create VendorWorkflowUtils instance for managing vendor workflows"""
+    from codemie_test_harness.tests.utils.vendor_workflow_utils import (
+        VendorWorkflowUtils,
+    )
+
+    return VendorWorkflowUtils()
 
 
 def pytest_sessionfinish(session):

codemie_test_harness/tests/enums/model_types.py

@@ -37,6 +37,7 @@ class ModelTypes(str, Enum):
     CLAUDE_4_OPUS = "claude-4-opus"
     CLAUDE_4_1_OPUS = "claude-4-1-opus"
     CLAUDE_4_SONNET_1M = "claude-4-sonnet-1m"
+    CLAUDE_4_5_HAIKU = "claude-4-5-haiku"
 
     # Other
     RLAB_QWQ_32B = "rlab-qwq-32b"

codemie_test_harness/tests/test_data/llm_test_data.py

@@ -56,6 +56,7 @@ MODEL_RESPONSES = [
     LlmResponseData(ModelTypes.CLAUDE_4_1_OPUS, AWS_ENVS),
     LlmResponseData(ModelTypes.CLAUDE_4_SONNET_1M, AWS_ENVS),
     LlmResponseData(ModelTypes.CLAUDE_4_5_SONNET, AWS_ENVS),
+    LlmResponseData(ModelTypes.CLAUDE_4_5_HAIKU, AWS_ENVS),
     # Other LLMs test data
     LlmResponseData(ModelTypes.RLAB_QWQ_32B, OTHER_ENVS),
     LlmResponseData(ModelTypes.DEEPSEEK_R1, OTHER_ENVS),

codemie_test_harness/tests/test_data/vendor_workflow_test_data.py

@@ -0,0 +1,50 @@
+"""Test data for vendor workflow e2e tests."""
+
+import json
+import pytest
+from codemie_sdk.models.vendor_assistant import VendorType
+from codemie_sdk.models.integration import CredentialTypes
+
+from codemie_test_harness.tests.utils.credentials_manager import CredentialsManager
+
+
+vendor_workflow_test_data = [
+    pytest.param(
+        VendorType.AWS,
+        CredentialTypes.AWS,
+        CredentialsManager.aws_credentials(),
+        "codemie-autotests-flow-object-input",
+        json.dumps({"genre": "pop", "number": 3}),
+        """
+            # Pop Playlist: Just Three Hits
+
+            1. "As It Was" by Harry Styles
+            2. "Blinding Lights" by The Weeknd
+            3. "Levitating" by Dua Lipa
+
+            These three modern pop classics offer a perfect mini-playlist with upbeat tempos and catchy hooks. Enjoy!
+        """,
+        marks=[pytest.mark.aws, pytest.mark.bedrock],
+        id="AWS_Bedrock_Object_Input",
+    ),
+    # pytest.param(
+    #     VendorType.AZURE,
+    #     CredentialTypes.AZURE,
+    #     CredentialsManager.azure_credentials(),
+    #     "workflow-name",
+    #     json.dumps({"genre": "rock", "number": 5}),
+    #     "Rock Playlist",
+    #     marks=[pytest.mark.azure],
+    #     id="Azure_AI_Object_Input",
+    # ),
+    # pytest.param(
+    #     VendorType.GCP,
+    #     CredentialTypes.GCP,
+    #     CredentialsManager.gcp_credentials(),
+    #     "workflow-name",
+    #     json.dumps({"genre": "jazz", "number": 4}),
+    #     "Jazz Playlist",
+    #     marks=[pytest.mark.gcp],
+    #     id="GCP_Vertex_AI_Object_Input",
+    # ),
+]

codemie_test_harness/tests/utils/vendor_utils.py

@@ -20,13 +20,13 @@ logger = setup_logger(__name__)
 DRAFT_VERSION = "DRAFT"
 
 
-class VendorUtils:
+class VendorAssistantUtils:
     """Utility class for vendor assistant operations."""
 
     def __init__(self):
         """Initialize VendorUtils with CodeMie client."""
         self.client = get_client()
-        self.vendor_service = self.client.vendors
+        self.vendor_assistant_service = self.client.vendor_assistants
 
     def get_assistant_settings(
         self,
@@ -47,7 +47,7 @@ class VendorUtils:
         logger.info(
             f"Getting assistant settings for {vendor.value} (page={page}, per_page={per_page})"
         )
-        settings = self.vendor_service.get_assistant_settings(
+        settings = self.vendor_assistant_service.get_assistant_settings(
             vendor=vendor, page=page, per_page=per_page
         )
         logger.info(
@@ -76,7 +76,7 @@ class VendorUtils:
         logger.info(
             f"Getting assistants for {vendor.value} setting {setting_id} (per_page={per_page})"
         )
-        assistants = self.vendor_service.get_assistants(
+        assistants = self.vendor_assistant_service.get_assistants(
             vendor=vendor,
             setting_id=setting_id,
             per_page=per_page,
@@ -106,7 +106,7 @@ class VendorUtils:
         logger.info(
             f"Getting assistant {assistant_id} for {vendor.value} setting {setting_id}"
         )
-        assistant = self.vendor_service.get_assistant(
+        assistant = self.vendor_assistant_service.get_assistant(
             vendor=vendor, assistant_id=assistant_id, setting_id=setting_id
         )
         logger.info(
@@ -137,7 +137,7 @@ class VendorUtils:
         logger.info(
             f"Getting aliases for assistant {assistant_id} in {vendor.value} setting {setting_id}"
         )
-        aliases = self.vendor_service.get_assistant_aliases(
+        aliases = self.vendor_assistant_service.get_assistant_aliases(
             vendor=vendor,
             assistant_id=assistant_id,
             setting_id=setting_id,
@@ -164,7 +164,7 @@ class VendorUtils:
             VendorAssistantInstallResponse containing installation summary with AI run IDs
         """
         logger.info(f"Installing {len(assistants)} assistant(s) for {vendor.value}")
-        response = self.vendor_service.install_assistants(
+        response = self.vendor_assistant_service.install_assistants(
             vendor=vendor, assistants=assistants
         )
         for item in response.summary:
@@ -188,7 +188,7 @@ class VendorUtils:
             VendorAssistantUninstallResponse with success status
         """
         logger.info(f"Uninstalling assistant with CodeMie ID: {codemie_id}")
-        response = self.vendor_service.uninstall_assistant(
+        response = self.vendor_assistant_service.uninstall_assistant(
             vendor=vendor, ai_run_id=codemie_id
         )
         if response.success:

codemie_test_harness/tests/utils/vendor_workflow_utils.py

@@ -0,0 +1,418 @@
+"""Utility class for managing vendor workflows (AWS Bedrock, Azure, GCP)."""
+
+from typing import Optional, List
+from codemie_sdk.models.vendor_assistant import VendorType
+from codemie_sdk.models.vendor_workflow import (
+    VendorWorkflowSettingsResponse,
+    VendorWorkflowsResponse,
+    VendorWorkflow,
+    VendorWorkflowAliasesResponse,
+    VendorWorkflowInstallRequest,
+    VendorWorkflowInstallResponse,
+    VendorWorkflowUninstallResponse,
+    VendorWorkflowStatus,
+)
+from codemie_test_harness.tests.utils.client_factory import get_client
+from codemie_test_harness.tests.utils.logger_util import setup_logger
+
+logger = setup_logger(__name__)
+
+DRAFT_VERSION = "DRAFT"
+
+
+class VendorWorkflowUtils:
+    """Utility class for vendor workflow operations."""
+
+    def __init__(self):
+        """Initialize VendorWorkflowUtils with CodeMie client."""
+        self.client = get_client()
+        self.vendor_workflow_service = self.client.vendor_workflows
+
+    def get_workflow_settings(
+        self,
+        vendor: VendorType,
+        page: int = 0,
+        per_page: int = 10,
+    ) -> VendorWorkflowSettingsResponse:
+        """Get workflow settings for a specific cloud vendor.
+
+        Args:
+            vendor: Cloud vendor type (AWS, AZURE, GCP)
+            page: Page number for pagination
+            per_page: Number of items per page
+
+        Returns:
+            VendorWorkflowSettingsResponse containing list of settings
+        """
+        logger.info(
+            f"Getting workflow settings for {vendor.value} (page={page}, per_page={per_page})"
+        )
+        settings = self.vendor_workflow_service.get_workflow_settings(
+            vendor=vendor, page=page, per_page=per_page
+        )
+        logger.info(
+            f"Retrieved {len(settings.data)} settings for {vendor.value} (total: {settings.pagination.total})"
+        )
+        return settings
+
+    def get_workflows(
+        self,
+        vendor: VendorType,
+        setting_id: str,
+        per_page: int = 10,
+        next_token: Optional[str] = None,
+    ) -> VendorWorkflowsResponse:
+        """Get workflows for a specific vendor setting.
+
+        Args:
+            vendor: Cloud vendor type
+            setting_id: ID of the vendor setting
+            per_page: Number of items per page
+            next_token: Token for pagination
+
+        Returns:
+            VendorWorkflowsResponse containing list of workflows
+        """
+        logger.info(
+            f"Getting workflows for {vendor.value} setting {setting_id} (per_page={per_page})"
+        )
+        workflows = self.vendor_workflow_service.get_workflows(
+            vendor=vendor,
+            setting_id=setting_id,
+            per_page=per_page,
+            next_token=next_token,
+        )
+        logger.info(
+            f"Retrieved {len(workflows.data)} workflows for setting {setting_id}"
+        )
+        return workflows
+
+    def get_workflow(
+        self,
+        vendor: VendorType,
+        workflow_id: str,
+        setting_id: str,
+    ) -> VendorWorkflow:
+        """Get a specific workflow by ID.
+
+        Args:
+            vendor: Cloud vendor type
+            workflow_id: ID of the workflow
+            setting_id: ID of the vendor setting
+
+        Returns:
+            VendorWorkflow with workflow details
+        """
+        logger.info(
+            f"Getting workflow {workflow_id} for {vendor.value} setting {setting_id}"
+        )
+        workflow = self.vendor_workflow_service.get_workflow(
+            vendor=vendor, workflow_id=workflow_id, setting_id=setting_id
+        )
+        logger.info(f"Retrieved workflow: {workflow.name} (status: {workflow.status})")
+        return workflow
+
+    def get_workflow_aliases(
+        self,
+        vendor: VendorType,
+        workflow_id: str,
+        setting_id: str,
+        per_page: int = 10,
+        next_token: Optional[str] = None,
+    ) -> VendorWorkflowAliasesResponse:
+        """Get aliases for a specific vendor workflow.
+
+        Args:
+            vendor: Cloud vendor type
+            workflow_id: ID of the workflow
+            setting_id: ID of the vendor setting
+            per_page: Number of items per page
+            next_token: Token for pagination
+
+        Returns:
+            VendorWorkflowAliasesResponse containing list of aliases
+        """
+        logger.info(
+            f"Getting aliases for workflow {workflow_id} in {vendor.value} setting {setting_id}"
+        )
+        aliases = self.vendor_workflow_service.get_workflow_aliases(
+            vendor=vendor,
+            workflow_id=workflow_id,
+            setting_id=setting_id,
+            per_page=per_page,
+            next_token=next_token,
+        )
+        logger.info(f"Retrieved {len(aliases.data)} aliases for workflow {workflow_id}")
+        return aliases
+
+    def install_workflows(
+        self,
+        vendor: VendorType,
+        workflows: List[VendorWorkflowInstallRequest],
+    ) -> VendorWorkflowInstallResponse:
+        """Install/activate vendor workflows.
+
+        Args:
+            vendor: Cloud vendor type
+            workflows: List of workflow installation requests
+
+        Returns:
+            VendorWorkflowInstallResponse containing installation summary with AI run IDs
+        """
+        logger.info(f"Installing {len(workflows)} workflow(s) for {vendor.value}")
+        response = self.vendor_workflow_service.install_workflows(
+            vendor=vendor, workflows=workflows
+        )
+        for item in response.summary:
+            logger.info(
+                f"Installed workflow {item.flowId} (alias: {item.flowAliasId}) -> CodeMie ID: {item.aiRunId}"
+            )
+        return response
+
+    def uninstall_workflow(
+        self,
+        vendor: VendorType,
+        codemie_id: str,
+    ) -> VendorWorkflowUninstallResponse:
+        """Uninstall/deactivate a vendor workflow.
+
+        Args:
+            vendor: Cloud vendor type
+            codemie_id: CodeMie workflow ID from installation
+
+        Returns:
+            VendorWorkflowUninstallResponse with success status
+        """
+        logger.info(f"Uninstalling workflow with CodeMie ID: {codemie_id}")
+        response = self.vendor_workflow_service.uninstall_workflow(
+            vendor=vendor, ai_run_id=codemie_id
+        )
+        if response.success:
+            logger.info(f"Successfully uninstalled workflow {codemie_id}")
+        else:
+            logger.warning(f"Failed to uninstall workflow {codemie_id}")
+        return response
+
+    def get_prepared_workflow(
+        self,
+        workflows: List[VendorWorkflow],
+    ) -> Optional[VendorWorkflow]:
+        """Get first PREPARED workflow from the list.
+
+        Args:
+            workflows: List of vendor workflows
+
+        Returns:
+            First PREPARED workflow or first workflow if none are PREPARED
+        """
+        return next(
+            (w for w in workflows if w.status == VendorWorkflowStatus.PREPARED),
+            workflows[0] if workflows else None,
+        )
+
+    def get_non_draft_alias(
+        self,
+        aliases: List,
+    ) -> Optional:
+        """Get first non-DRAFT alias from the list.
+
+        Args:
+            aliases: List of vendor workflow aliases
+
+        Returns:
+            First non-DRAFT alias or first alias if all are DRAFT
+        """
+        return next(
+            (a for a in aliases if a.version != DRAFT_VERSION),
+            aliases[0] if aliases else None,
+        )
+
+    def find_workflow_by_name(
+        self,
+        vendor: VendorType,
+        setting_id: str,
+        workflow_name: str,
+    ) -> Optional[tuple[VendorWorkflow, str]]:
+        """Find a workflow by name and return it with an alias.
+
+        Args:
+            vendor: Cloud vendor type
+            setting_id: ID of the vendor setting
+            workflow_name: Name of the workflow to find
+
+        Returns:
+            Tuple of (VendorWorkflow, alias_id) or None if workflow not found
+        """
+        logger.info(
+            f"Searching for workflow '{workflow_name}' in {vendor.value} setting {setting_id}"
+        )
+        workflows_response = self.get_workflows(vendor=vendor, setting_id=setting_id)
+
+        for workflow in workflows_response.data:
+            if (
+                workflow.name == workflow_name
+                and workflow.status == VendorWorkflowStatus.PREPARED
+            ):
+                aliases_response = self.get_workflow_aliases(
+                    vendor=vendor, workflow_id=workflow.id, setting_id=setting_id
+                )
+                if aliases_response.data:
+                    non_draft_alias = self.get_non_draft_alias(aliases_response.data)
+                    if non_draft_alias:
+                        logger.info(
+                            f"Found workflow: {workflow.name} (ID: {workflow.id}, Alias: {non_draft_alias.id})"
+                        )
+                        return workflow, non_draft_alias.id
+
+        logger.warning(
+            f"Workflow '{workflow_name}' not found for {vendor.value} setting {setting_id}"
+        )
+        return None
+
+    def find_first_available_workflow(
+        self,
+        vendor: VendorType,
+        setting_id: str,
+    ) -> Optional[tuple[VendorWorkflow, str]]:
+        """Find the first available (PREPARED) workflow with an alias.
+
+        Args:
+            vendor: Cloud vendor type
+            setting_id: ID of the vendor setting
+
+        Returns:
+            Tuple of (VendorWorkflow, alias_id) or None if no available workflow found
+        """
+        logger.info(
+            f"Searching for available workflow in {vendor.value} setting {setting_id}"
+        )
+        workflows_response = self.get_workflows(vendor=vendor, setting_id=setting_id)
+
+        for workflow in workflows_response.data:
+            if workflow.status == VendorWorkflowStatus.PREPARED:
+                aliases_response = self.get_workflow_aliases(
+                    vendor=vendor, workflow_id=workflow.id, setting_id=setting_id
+                )
+                if aliases_response.data:
+                    first_alias = aliases_response.data[0]
+                    logger.info(
+                        f"Found available workflow: {workflow.name} (ID: {workflow.id}, Alias: {first_alias.id})"
+                    )
+                    return workflow, first_alias.id
+
+        logger.warning(
+            f"No available workflow found for {vendor.value} setting {setting_id}"
+        )
+        return None
+
+    def install_workflow_by_name(
+        self,
+        vendor: VendorType,
+        setting_id: str,
+        workflow_name: str,
+    ) -> Optional[str]:
+        """Find and install a workflow by name.
+
+        Args:
+            vendor: Cloud vendor type
+            setting_id: ID of the vendor setting
+            workflow_name: Name of the workflow to install
+
+        Returns:
+            CodeMie ID of the installed workflow or None if workflow not found
+        """
+        result = self.find_workflow_by_name(
+            vendor=vendor, setting_id=setting_id, workflow_name=workflow_name
+        )
+        if not result:
+            return None
+
+        workflow, alias_id = result
+        install_request = VendorWorkflowInstallRequest(
+            id=workflow.id,
+            flowAliasId=alias_id,
+            setting_id=setting_id,
+        )
+
+        install_response = self.install_workflows(
+            vendor=vendor, workflows=[install_request]
+        )
+
+        if install_response.summary:
+            return install_response.summary[0].aiRunId
+
+        return None
+
+    def install_first_available_workflow(
+        self,
+        vendor: VendorType,
+        setting_id: str,
+    ) -> Optional[str]:
+        """Find and install the first available workflow.
+
+        Args:
+            vendor: Cloud vendor type
+            setting_id: ID of the vendor setting
+
+        Returns:
+            CodeMie ID of the installed workflow or None if no workflow available
+        """
+        result = self.find_first_available_workflow(
+            vendor=vendor, setting_id=setting_id
+        )
+        if not result:
+            return None
+
+        workflow, alias_id = result
+        install_request = VendorWorkflowInstallRequest(
+            id=workflow.id,
+            flowAliasId=alias_id,
+            setting_id=setting_id,
+        )
+
+        install_response = self.install_workflows(
+            vendor=vendor, workflows=[install_request]
+        )
+
+        if install_response.summary:
+            return install_response.summary[0].aiRunId
+
+        return None
+
+    def find_setting_for_integration(
+        self,
+        vendor: VendorType,
+        integration_id: str,
+    ):
+        """Find setting for an integration by paginating through all settings.
+
+        Args:
+            vendor: Type of vendor (AWS, AZURE, GCP)
+            integration_id: ID of the integration to find (searches by setting_id)
+
+        Returns:
+            VendorWorkflowSetting or None if not found
+        """
+        page = 0
+        per_page = 50
+
+        while True:
+            settings_response = self.get_workflow_settings(
+                vendor=vendor,
+                page=page,
+                per_page=per_page,
+            )
+
+            # Find the setting for our integration by setting_id
+            for s in settings_response.data:
+                if s.setting_id == integration_id:
+                    return s
+
+            # Check if there are more pages
+            if page >= settings_response.pagination.pages - 1:
+                break
+
+            page += 1
+
+        logger.warning(f"Setting not found for integration ID '{integration_id}'")
+        return None

codemie_test_harness/tests/vendor/assistants/test_vendor_assistants.py

@@ -12,7 +12,7 @@ from codemie_test_harness.tests.test_data.vendor_test_data import (
     vendor_assistant_test_data,
 )
 def test_vendor_assistant_installation_and_chat(
-    vendor_utils,
+    vendor_assistant_utils,
     integration,
     assistant_utils,
     similarity_check,
@@ -25,21 +25,21 @@ def test_vendor_assistant_installation_and_chat(
     """Test vendor assistant installation and chat functionality."""
     _integration = integration(credential_type, credentials)
 
-    setting = vendor_utils.find_setting_for_integration(
+    setting = vendor_assistant_utils.find_setting_for_integration(
         vendor=vendor_type,
         integration_id=_integration.id,
     )
     assert_that(setting, is_not(none()))
     assert_that(setting.invalid or False, is_(False))
 
-    assistants_response = vendor_utils.get_assistants(
+    assistants_response = vendor_assistant_utils.get_assistants(
         vendor=vendor_type,
         setting_id=setting.setting_id,
     )
     assert_that(assistants_response, is_not(none()))
     assert_that(assistants_response.data, is_not(empty()))
 
-    codemie_id = vendor_utils.install_first_available_assistant(
+    codemie_id = vendor_assistant_utils.install_first_available_assistant(
         vendor=vendor_type,
         setting_id=setting.setting_id,
     )
@@ -58,7 +58,7 @@ def test_vendor_assistant_installation_and_chat(
 
     similarity_check.check_similarity(response, expected_response)
 
-    uninstall_response = vendor_utils.uninstall_assistant(
+    uninstall_response = vendor_assistant_utils.uninstall_assistant(
         vendor=vendor_type,
         codemie_id=codemie_id,
     )