codemie-test-harness 0.1.159__py3-none-any.whl → 0.1.160__py3-none-any.whl

codemie_test_harness/tests/integrations/user/test_user_integrations.py

@@ -7,7 +7,6 @@ from codemie_sdk.models.integration import (
     Integration,
     IntegrationTestRequest,
     IntegrationTestResponse,
-    CredentialTypes,
 )
 from codemie_test_harness.tests import PROJECT
 from codemie_test_harness.tests.test_data.integrations_test_data import (
@@ -28,31 +27,6 @@ from codemie_test_harness.tests.utils.base_utils import (
 @pytest.mark.parametrize(
     "credential_type, credentials",
     valid_integrations,
-    ids=[
-        f"User integration: {CredentialTypes.AWS}",
-        f"User integration: {CredentialTypes.AZURE}",
-        f"User integration: {CredentialTypes.GCP}",
-        f"User integration: {CredentialTypes.SONAR}",
-        f"User integration: {CredentialTypes.SONAR}",
-        f"User integration: {CredentialTypes.GIT}",
-        f"User integration: {CredentialTypes.GIT}",
-        f"User integration: {CredentialTypes.CONFLUENCE}",
-        f"User integration: {CredentialTypes.CONFLUENCE}",
-        f"User integration: {CredentialTypes.JIRA}",
-        f"User integration: {CredentialTypes.JIRA}",
-        f"User integration: {CredentialTypes.SQL}",
-        f"User integration: {CredentialTypes.SQL}",
-        f"User integration: {CredentialTypes.ELASTIC}",
-        f"User integration: {CredentialTypes.MCP}",
-        f"User integration: {CredentialTypes.AZURE_DEVOPS}",
-        f"User integration: {CredentialTypes.FILESYSTEM}",
-        f"User integration: {CredentialTypes.EMAIL}",
-        f"User integration: {CredentialTypes.TELEGRAM}",
-        f"User integration: {CredentialTypes.SERVICE_NOW}",
-        f"User integration: {CredentialTypes.KUBERNETES}",
-        f"User integration: {CredentialTypes.KEYCLOAK}",
-        f"User integration: {CredentialTypes.REPORT_PORTAL}",
-    ],
 )
 def test_integration_created(
     credential_type, credentials, general_integration, integration_utils
@@ -86,21 +60,6 @@ def test_integration_created(
 @pytest.mark.parametrize(
     "credential_type, credentials",
     testable_integrations,
-    ids=[
-        f"User integration: {CredentialTypes.AWS}",
-        f"User integration: {CredentialTypes.AZURE}",
-        f"User integration: {CredentialTypes.GCP}",
-        f"User integration: {CredentialTypes.SONAR}",
-        f"User integration: {CredentialTypes.SONAR}",
-        f"User integration: {CredentialTypes.CONFLUENCE}",
-        f"User integration: {CredentialTypes.CONFLUENCE}",
-        f"User integration: {CredentialTypes.JIRA}",
-        f"User integration: {CredentialTypes.JIRA}",
-        f"User integration: {CredentialTypes.EMAIL}",
-        f"User integration: {CredentialTypes.SERVICE_NOW}",
-        f"User integration: {CredentialTypes.KUBERNETES}",
-        f"User integration: {CredentialTypes.REPORT_PORTAL}",
-    ],
 )
 def test_integration_after_creation(
     credential_type, credentials, general_integration, integration_utils
@@ -139,31 +98,6 @@ def test_integration_after_creation(
 @pytest.mark.parametrize(
     "credential_type, credentials",
     valid_integrations,
-    ids=[
-        f"User integration: {CredentialTypes.AWS}",
-        f"User integration: {CredentialTypes.AZURE}",
-        f"User integration: {CredentialTypes.GCP}",
-        f"User integration: {CredentialTypes.SONAR}",
-        f"User integration: {CredentialTypes.SONAR}",
-        f"User integration: {CredentialTypes.GIT}",
-        f"User integration: {CredentialTypes.GIT}",
-        f"User integration: {CredentialTypes.CONFLUENCE}",
-        f"User integration: {CredentialTypes.CONFLUENCE}",
-        f"User integration: {CredentialTypes.JIRA}",
-        f"User integration: {CredentialTypes.JIRA}",
-        f"User integration: {CredentialTypes.SQL}",
-        f"User integration: {CredentialTypes.SQL}",
-        f"User integration: {CredentialTypes.ELASTIC}",
-        f"User integration: {CredentialTypes.MCP}",
-        f"User integration: {CredentialTypes.AZURE_DEVOPS}",
-        f"User integration: {CredentialTypes.FILESYSTEM}",
-        f"User integration: {CredentialTypes.EMAIL}",
-        f"User integration: {CredentialTypes.TELEGRAM}",
-        f"User integration: {CredentialTypes.SERVICE_NOW}",
-        f"User integration: {CredentialTypes.KUBERNETES}",
-        f"User integration: {CredentialTypes.KEYCLOAK}",
-        f"User integration: {CredentialTypes.REPORT_PORTAL}",
-    ],
 )
 def test_update_integration(
     credential_type, credentials, general_integration, integration_utils
@@ -209,19 +143,6 @@ def test_update_integration(
 @pytest.mark.parametrize(
     "credential_type, credentials, error_message",
     invalid_integrations,
-    ids=[
-        f"User integration: {CredentialTypes.AWS}",
-        f"User integration: {CredentialTypes.AZURE}",
-        f"User integration: {CredentialTypes.GCP}",
-        f"User integration: {CredentialTypes.SONAR}",
-        f"User integration: {CredentialTypes.SONAR}",
-        f"User integration: {CredentialTypes.EMAIL}",
-        f"User integration: {CredentialTypes.JIRA}",
-        f"User integration: {CredentialTypes.CONFLUENCE}",
-        f"User integration: {CredentialTypes.SERVICE_NOW}",
-        f"User integration: {CredentialTypes.KUBERNETES}",
-        f"User integration: {CredentialTypes.REPORT_PORTAL}",
-    ],
 )
 def test_integration_with_invalid_credentials(
     credential_type, credentials, error_message, integration_utils
@@ -263,21 +184,6 @@ def test_integration_with_invalid_credentials(
 @pytest.mark.parametrize(
     "credential_type, credentials",
     testable_integrations,
-    ids=[
-        f"User integration: {CredentialTypes.AWS}",
-        f"User integration: {CredentialTypes.AZURE}",
-        f"User integration: {CredentialTypes.GCP}",
-        f"User integration: {CredentialTypes.SONAR}",
-        f"User integration: {CredentialTypes.SONAR}",
-        f"User integration: {CredentialTypes.CONFLUENCE}",
-        f"User integration: {CredentialTypes.CONFLUENCE}",
-        f"User integration: {CredentialTypes.JIRA}",
-        f"User integration: {CredentialTypes.JIRA}",
-        f"User integration: {CredentialTypes.EMAIL}",
-        f"User integration: {CredentialTypes.SERVICE_NOW}",
-        f"User integration: {CredentialTypes.KUBERNETES}",
-        f"User integration: {CredentialTypes.REPORT_PORTAL}",
-    ],
 )
 def test_integration_during_creation(credential_type, credentials, integration_utils):
     integration_test_model = IntegrationTestRequest(
@@ -302,19 +208,6 @@ def test_integration_during_creation(credential_type, credentials, integration_u
 @pytest.mark.parametrize(
     "credential_type, credentials, error_message",
     invalid_integrations,
-    ids=[
-        f"User integration: {CredentialTypes.AWS}",
-        f"User integration: {CredentialTypes.AZURE}",
-        f"User integration: {CredentialTypes.GCP}",
-        f"User integration: {CredentialTypes.SONAR}",
-        f"User integration: {CredentialTypes.SONAR}",
-        f"User integration: {CredentialTypes.EMAIL}",
-        f"User integration: {CredentialTypes.JIRA}",
-        f"User integration: {CredentialTypes.CONFLUENCE}",
-        f"User integration: {CredentialTypes.SERVICE_NOW}",
-        f"User integration: {CredentialTypes.KUBERNETES}",
-        f"User integration: {CredentialTypes.REPORT_PORTAL}",
-    ],
 )
 def test_integration_during_creation_with_invalid_credentials(
     credential_type, credentials, error_message, integration_utils
@@ -338,31 +231,6 @@ def test_integration_during_creation_with_invalid_credentials(
 @pytest.mark.parametrize(
     "credential_type, credentials",
     valid_integrations,
-    ids=[
-        f"User integration: {CredentialTypes.AWS}",
-        f"User integration: {CredentialTypes.AZURE}",
-        f"User integration: {CredentialTypes.GCP}",
-        f"User integration: {CredentialTypes.SONAR}",
-        f"User integration: {CredentialTypes.SONAR}",
-        f"User integration: {CredentialTypes.GIT}",
-        f"User integration: {CredentialTypes.GIT}",
-        f"User integration: {CredentialTypes.CONFLUENCE}",
-        f"User integration: {CredentialTypes.CONFLUENCE}",
-        f"User integration: {CredentialTypes.JIRA}",
-        f"User integration: {CredentialTypes.JIRA}",
-        f"User integration: {CredentialTypes.SQL}",
-        f"User integration: {CredentialTypes.SQL}",
-        f"User integration: {CredentialTypes.ELASTIC}",
-        f"User integration: {CredentialTypes.MCP}",
-        f"User integration: {CredentialTypes.AZURE_DEVOPS}",
-        f"User integration: {CredentialTypes.FILESYSTEM}",
-        f"User integration: {CredentialTypes.EMAIL}",
-        f"User integration: {CredentialTypes.TELEGRAM}",
-        f"User integration: {CredentialTypes.SERVICE_NOW}",
-        f"User integration: {CredentialTypes.KUBERNETES}",
-        f"User integration: {CredentialTypes.KEYCLOAK}",
-        f"User integration: {CredentialTypes.REPORT_PORTAL}",
-    ],
 )
 def test_delete_integration(
     credential_type, credentials, general_integration, integration_utils

codemie_test_harness/tests/llm/assistants/test_llm.py

@@ -1,9 +1,9 @@
-import os
 import pytest
 from hamcrest import assert_that, has_item
 from codemie_test_harness.tests.enums.model_types import ModelTypes
 from codemie_test_harness.tests.test_data.llm_test_data import MODEL_RESPONSES
 from codemie_test_harness.tests.utils.client_factory import get_client
+from codemie_test_harness.tests.utils.env_resolver import get_environment
 from codemie_test_harness.tests.utils.pytest_utils import check_mark
 
 
@@ -11,7 +11,7 @@ def pytest_generate_tests(metafunc):
     if "model_type" in metafunc.fixturenames:
         is_smoke = check_mark(metafunc, "smoke")
         test_data = []
-        env = os.getenv("ENV")
+        env = get_environment()
         if is_smoke:
             available_models = get_client().llms.list()
             for model in available_models:
@@ -23,7 +23,7 @@ def pytest_generate_tests(metafunc):
                         model_data.model_type,
                         marks=pytest.mark.skipif(
                             env not in model_data.environments,
-                            reason=f"Skip on non {'/'.join(model_data.environments[:-1])} envs",
+                            reason=f"Skip on non {'/'.join(str(env) for env in model_data.environments[:-1])} envs",
                         ),
                     )
                 )

codemie_test_harness/tests/service/test_assistant_service.py

@@ -1,4 +1,3 @@
-import os
 import uuid
 from datetime import datetime
 
@@ -34,6 +33,7 @@ from hamcrest import (
 
 from codemie_test_harness.tests import PROJECT
 from codemie_test_harness.tests.utils.base_utils import get_random_name
+from codemie_test_harness.tests.utils.env_resolver import EnvironmentResolver
 
 
 def validate_assistant_full_response(assistant):
@@ -498,7 +498,7 @@ def test_get_prebuilt_assistants(assistant_utils):
 
 
 @pytest.mark.skipif(
-    os.getenv("ENV") != "preview", reason="valid_assistant_id is for preview env"
+    not EnvironmentResolver.is_preview(), reason="valid_assistant_id is for preview env"
 )
 def test_assistant_evaluate(assistant_utils):
     valid_assistant_id = "05959338-06de-477d-9cc3-08369f858057"

codemie_test_harness/tests/test_data/cloud_tools_test_data.py

@@ -1,10 +1,13 @@
-import os
-
 import pytest
 
 from codemie_test_harness.tests.enums.tools import Toolkit, CloudTool
 from codemie_sdk.models.integration import CredentialTypes
 from codemie_test_harness.tests.utils.aws_parameters_store import CredentialsUtil
+from codemie_test_harness.tests.utils.env_resolver import (
+    EnvironmentResolver,
+    get_environment,
+)
+from codemie_test_harness.tests.enums.environment import Environment
 
 cloud_test_data = [
     pytest.param(
@@ -17,15 +20,30 @@ cloud_test_data = [
             Query example: {'query': {'service': 's3', 'method_name': 'list_buckets', 'method_arguments': {}}}
         """,
         """
-            Okay, here are the names of your S3 buckets:
-
-            *   codemie-bucket
-            *   codemie-terraform-states-025066278959
-            *   epam-cloud-s3-access-logs-025066278959-eu-central-1
-            *   epam-cloud-s3-access-logs-025066278959-us-east-1
-            *   terraform-states-025066278959
+            Here are the names of the S3 buckets:
+            
+            - az-codemie-terraform-states-025066278959
+            - codemie-bucket
+            - codemie-terraform-states-025066278959
+            - codemie-terraform-states-da-025066278959
+            - codemie-terraform-states-dar-025066278959
+            - codemie-terraform-states-ihorb-025066278959
+            - codemie-terraform-states-sh-test-025066278959
+            - da-codemie-user-data-025066278959
+            - epam-cloud-s3-access-logs-025066278959-eu-central-1
+            - epam-cloud-s3-access-logs-025066278959-eu-north-1
+            - epam-cloud-s3-access-logs-025066278959-eu-west-2
+            - epam-cloud-s3-access-logs-025066278959-us-east-1
+            - ggg-codemie-terraform-states-025066278959
+            - gggcodemie-user-data-025066278959
+            - pc-codemie-terraform-states-025066278959
+            - pc-codemie-user-data-025066278959
+            - sk-codemie-terraform-states-025066278959
+            - sk-codemie-user-data-025066278959
+            - terraform-states-025066278959
         """,
         marks=pytest.mark.aws,
+        id=CredentialTypes.AWS,
     ),
     pytest.param(
         Toolkit.CLOUD,
@@ -45,6 +63,7 @@ cloud_test_data = [
             Provisioning State: Succeeded
         """,
         marks=pytest.mark.azure,
+        id=CredentialTypes.AZURE,
     ),
     pytest.param(
         Toolkit.CLOUD,
@@ -106,10 +125,11 @@ cloud_test_data = [
         marks=[
             pytest.mark.gcp,
             pytest.mark.skipif(
-                os.getenv("ENV") in ["azure", "gcp"],
+                get_environment() in [Environment.AZURE, Environment.GCP],
                 reason="Still have an issue with encoding long strings",
             ),
         ],
+        id=CredentialTypes.GCP,
     ),
     pytest.param(
         Toolkit.CLOUD,
@@ -134,9 +154,10 @@ cloud_test_data = [
         marks=[
             pytest.mark.kubernetes,
             pytest.mark.skipif(
-                os.getenv("ENV") == "azure",
+                EnvironmentResolver.is_azure(),
                 reason="Still have an issue with encoding long strings",
             ),
         ],
+        id=CredentialTypes.KUBERNETES,
     ),
 ]

codemie_test_harness/tests/test_data/codebase_tools_test_data.py

@@ -57,6 +57,7 @@ sonar_tools_test_data = [
         """,
         "Yes, I have an access to SonarQube",
         marks=pytest.mark.sonar,
+        id=CodeBaseTool.SONAR,
     ),
     pytest.param(
         Toolkit.CODEBASE_TOOLS,
@@ -71,5 +72,6 @@ sonar_tools_test_data = [
         """,
         "Yes, I have an access to SonarCloud",
         marks=pytest.mark.sonar,
+        id=CodeBaseTool.SONAR_CLOUD,
     ),
 ]

codemie_test_harness/tests/test_data/data_management_tools_test_data.py

@@ -1,8 +1,7 @@
-import os
-
 import pytest
 
 from codemie_test_harness.tests.enums.integrations import DataBaseDialect
+from codemie_test_harness.tests.utils.env_resolver import EnvironmentResolver
 
 ELASTIC_TOOL_TASK = """
     send the query to Elastic:
@@ -59,8 +58,9 @@ sql_tools_test_data = [
     pytest.param(
         DataBaseDialect.MS_SQL,
         marks=pytest.mark.skipif(
-            os.getenv("ENV") not in ("aws", "azure", "gcp"),
+            not EnvironmentResolver.is_sandbox(),
             reason="MS SQL is only available in staging environments",
         ),
+        id=DataBaseDialect.MS_SQL,
     ),
 ]

codemie_test_harness/tests/test_data/direct_tools/cloud_tools_test_data.py

@@ -1,10 +1,9 @@
-import os
-
 import pytest
 
 from codemie_test_harness.tests.enums.tools import Toolkit, CloudTool
 from codemie_sdk.models.integration import CredentialTypes
 from codemie_test_harness.tests.utils.aws_parameters_store import CredentialsUtil
+from codemie_test_harness.tests.utils.env_resolver import EnvironmentResolver
 
 cloud_test_data = [
     pytest.param(
@@ -100,6 +99,7 @@ cloud_test_data = [
             }
         """,
         marks=pytest.mark.aws,
+        id=CredentialTypes.AWS,
     ),
     pytest.param(
         Toolkit.CLOUD,
@@ -125,6 +125,7 @@ cloud_test_data = [
             }
         """,
         marks=pytest.mark.azure,
+        id=CredentialTypes.AZURE,
     ),
     pytest.param(
         Toolkit.CLOUD,
@@ -170,10 +171,11 @@ cloud_test_data = [
         marks=[
             pytest.mark.gcp,
             pytest.mark.skipif(
-                os.getenv("ENV") == "azure",
+                EnvironmentResolver.is_azure(),
                 reason="Still have an issue with encoding long strings",
             ),
         ],
+        id=CredentialTypes.GCP,
     ),
     pytest.param(
         Toolkit.CLOUD,
@@ -844,9 +846,10 @@ cloud_test_data = [
         marks=[
             pytest.mark.kubernetes,
             pytest.mark.skipif(
-                os.getenv("ENV") == "azure",
+                EnvironmentResolver.is_azure(),
                 reason="Still have an issue with encoding long strings",
             ),
         ],
+        id=CredentialTypes.KUBERNETES,
     ),
 ]

codemie_test_harness/tests/test_data/direct_tools/codebase_tools_test_data.py

@@ -77,6 +77,7 @@ sonar_tools_test_data = [
             }
         """,
         marks=pytest.mark.sonar,
+        id=CodeBaseTool.SONAR,
     ),
     pytest.param(
         Toolkit.CODEBASE_TOOLS,
@@ -160,5 +161,6 @@ sonar_tools_test_data = [
             }
         """,
         marks=pytest.mark.sonar,
+        id=CodeBaseTool.SONAR_CLOUD,
     ),
 ]

codemie_test_harness/tests/test_data/direct_tools/data_management_tools_test_data.py

@@ -1,9 +1,8 @@
-import os
-
 import pytest
 
-from codemie_test_harness.tests.enums.integrations import DataBaseDialect
 from codemie_test_harness.tests.enums.tools import DataManagementTool, Toolkit
+from codemie_test_harness.tests.enums.integrations import DataBaseDialect
+from codemie_test_harness.tests.utils.env_resolver import EnvironmentResolver
 
 sql_tools_test_data = [
     (
@@ -40,9 +39,9 @@ sql_tools_test_data = [
         },
         [{"table_name": "Users"}, {"table_name": "Products"}],
         marks=pytest.mark.skipif(
-            os.getenv("ENV") not in ("aws", "azure", "gcp"),
+            not EnvironmentResolver.is_sandbox(),
             reason="MS SQL is only available in staging environments",
         ),
-        id="DataBaseDialect.MS_SQL",
+        id=DataBaseDialect.MS_SQL,
     ),
 ]

codemie_test_harness/tests/test_data/direct_tools/file_management_tools_test_data.py

@@ -1,4 +1,4 @@
-import os
+from codemie_test_harness.tests.utils.env_resolver import EnvironmentResolver
 
 import pytest
 
@@ -41,7 +41,7 @@ file_management_tools_test_data = [
         codemie-ui 
         """,
         marks=pytest.mark.skipif(
-            os.getenv("ENV") == "local",
+            EnvironmentResolver.is_localhost(),
             reason="Skipping this test on local environment",
         ),
     ),

codemie_test_harness/tests/test_data/direct_tools/notification_tools_test_data.py

@@ -1,4 +1,5 @@
-import os
+from codemie_test_harness.tests.enums.environment import Environment
+from codemie_test_harness.tests.utils.env_resolver import get_environment
 
 import pytest
 
@@ -21,10 +22,11 @@ notification_tools_test_data = [
         marks=[
             pytest.mark.email,
             pytest.mark.skipif(
-                os.getenv("ENV") in ["local", "gcp"],
+                get_environment() in [Environment.LOCALHOST, Environment.GCP],
                 reason="Skipping this test on local environment",
             ),
         ],
+        id=NotificationTool.EMAIL,
     ),
     pytest.param(
         NotificationTool.TELEGRAM,
@@ -54,5 +56,6 @@ notification_tools_test_data = [
             },
         },
         marks=pytest.mark.telegram,
+        id=NotificationTool.TELEGRAM,
     ),
 ]

codemie_test_harness/tests/test_data/direct_tools/project_management_tools_test_data.py

@@ -158,6 +158,7 @@ project_management_tools_data = [
              _17713":null,"customfield_27500":null}}
         """,
         marks=pytest.mark.jira,
+        id=ProjectManagementIntegrationType.JIRA,
     ),
     pytest.param(
         Toolkit.PROJECT_MANAGEMENT,
@@ -188,5 +189,6 @@ project_management_tools_data = [
         expand=body.storage&cql=title+~+%22AQA+backlog+estimation%22","base":"https://kb.epam.com","context":""}}
         """,
         marks=pytest.mark.confluence,
+        id=ProjectManagementIntegrationType.CONFLUENCE,
     ),
 ]

codemie_test_harness/tests/test_data/direct_tools/research_tools_test_data.py

@@ -34,6 +34,7 @@ research_tools_test_data = [
         marks=pytest.mark.skip(
             reason="Temporarily skipping Tavily test until it is fixed"
         ),
+        id=ResearchToolName.TAVILY_SEARCH,
     ),
     (
         ResearchToolName.GOOGLE_PLACES,

codemie_test_harness/tests/test_data/direct_tools/vcs_tools_test_data.py

@@ -1,4 +1,5 @@
 import pytest
+from codemie_sdk.models.integration import CredentialTypes
 
 from codemie_test_harness.tests.enums.tools import VcsTool, Toolkit
 
@@ -91,6 +92,7 @@ vcs_tools_test_data = [
             }
         """,
         marks=pytest.mark.github,
+        id=f"{CredentialTypes.GIT}_github",
     ),
     pytest.param(
         Toolkit.VCS,
@@ -135,5 +137,6 @@ vcs_tools_test_data = [
         
         """,
         marks=pytest.mark.gitlab,
+        id=f"{CredentialTypes.GIT}_gitlab",
     ),
 ]

codemie_test_harness/tests/test_data/file_management_tools_test_data.py

@@ -1,4 +1,4 @@
-import os
+from codemie_test_harness.tests.utils.env_resolver import EnvironmentResolver
 
 import pytest
 
@@ -89,29 +89,33 @@ GENERATE_IMAGE_TOOL_TASK = """
 """
 
 file_management_tools_test_data = [
-    (
+    pytest.param(
         FileManagementTool.PYTHON_CODE_INTERPRETER,
         CODE_INTERPRETER_TOOL_TASK,
         RESPONSE_FOR_CODE_INTERPRETER,
+        id=FileManagementTool.PYTHON_CODE_INTERPRETER,
     ),
     pytest.param(
         FileManagementTool.LIST_DIRECTORY,
         LIST_DIR_TOOL_TASK,
         RESPONSE_FOR_LIST_DIR,
         marks=pytest.mark.skipif(
-            os.getenv("ENV") == "local",
+            EnvironmentResolver.is_localhost(),
             reason="Skipping this test on local environment",
         ),
+        id=FileManagementTool.LIST_DIRECTORY,
     ),
-    (
+    pytest.param(
         FileManagementTool.WRITE_FILE,
         WRITE_FILE_TASK,
         RESPONSE_FOR_WRITE_FILE_TASK,
+        id=FileManagementTool.WRITE_FILE,
     ),
-    (
+    pytest.param(
         FileManagementTool.RUN_COMMAND_LINE,
         COMMAND_LINE_TOOL_TASK,
         RESPONSE_FOR_COMMAND_LINE_TASK,
+        id=FileManagementTool.RUN_COMMAND_LINE,
     ),
 ]
 

codemie_test_harness/tests/test_data/index_test_data.py

@@ -1,32 +1,30 @@
-import os
 from dataclasses import dataclass
 from typing import List
 
 import pytest
 
+from codemie_test_harness.tests.utils.env_resolver import get_environment
+from codemie_test_harness.tests.enums.environment import Environment
+
 
 @dataclass
 class EmbeddingData:
     """Data class to store Embedding models."""
 
     model_type: str
-    environments: List[str]
-
+    environments: List[Environment]
 
-AZURE_ENVS = ["preview", "azure", "local"]
-GCP_ENVS = ["preview", "gcp", "local"]
-AWS_ENVS = ["preview", "aws", "local"]
 
 MODELS = [
-    EmbeddingData("titan", AWS_ENVS),
-    EmbeddingData("gecko", GCP_ENVS),
-    EmbeddingData("ada-002", AZURE_ENVS),
+    EmbeddingData("titan", Environment.get_aws_environments()),
+    EmbeddingData("gecko", Environment.get_gcp_environments()),
+    EmbeddingData("ada-002", Environment.get_azure_environments()),
 ]
 
 
 def generate_test_data():
     """Generate pytest parameters for Embedding models"""
-    env = os.getenv("ENV")
+    env = get_environment()
     test_data = []
 
     for model in MODELS:
@@ -35,7 +33,7 @@ def generate_test_data():
                 model.model_type,
                 marks=pytest.mark.skipif(
                     env not in model.environments,
-                    reason=f"Skip on non {'/'.join(model.environments[:-1])} envs",
+                    reason=f"Skip on non {'/'.join(str(env) for env in model.environments[:-1])} envs",
                 ),
             )
         )