codemie-test-harness 0.1.136__py3-none-any.whl → 0.1.137__py3-none-any.whl

codemie_test_harness/tests/service/test_datasource_service.py

@@ -1,23 +1,21 @@
-"""Integration tests for LLMService."""
-
 import os
-from time import sleep
-
 import pytest
+from hamcrest import (
+    assert_that,
+    is_not,
+    instance_of,
+    has_property,
+    has_length,
+    greater_than,
+    is_,
+    is_in,
+    all_of,
+    equal_to,
+)
 
-from codemie_sdk import CodeMieClient
 from codemie_sdk.models.datasource import (
     DataSourceType,
     DataSourceStatus,
-    CodeDataSourceRequest,
-    CodeDataSourceType,
-    UpdateCodeDataSourceRequest,
-    ConfluenceDataSourceRequest,
-    UpdateConfluenceDataSourceRequest,
-    UpdateJiraDataSourceRequest,
-    JiraDataSourceRequest,
-    GoogleDataSourceRequest,
-    UpdateGoogleDataSourceRequest,
     Jira,
     Confluence,
     Code,
@@ -25,24 +23,12 @@ from codemie_sdk.models.datasource import (
 from codemie_sdk.models.integration import (
     CredentialTypes,
     CredentialValues,
-    Integration,
     IntegrationType,
 )
 from codemie_test_harness.tests import PROJECT
 from codemie_test_harness.tests.utils.base_utils import get_random_name
 
 
-@pytest.fixture
-def project_name():
-    """Return project name for tests."""
-    return PROJECT
-
-
-@pytest.fixture
-def datasource_name():
-    return get_random_name()
-
-
 @pytest.fixture
 def integration_config():
     """Configuration for different integration types."""
@@ -68,72 +54,33 @@ def integration_config():
 class TestDatasourceBase:
     """Base class for datasource tests with common utility methods."""
 
-    @staticmethod
-    def create_integration(
-        client: CodeMieClient,
-        project_name: str,
-        cred_type: CredentialTypes,
-        config: dict,
-    ) -> Integration:
-        """Create integration with given credentials."""
-        integration_alias = get_random_name()
-        credential_values = [
-            CredentialValues(key=k, value=v) for k, v in config[cred_type].items()
-        ]
-
-        integration_request = Integration(
-            project_name=project_name,
-            credential_type=cred_type,
-            credential_values=credential_values,
-            alias=integration_alias,
-            setting_type=IntegrationType.USER,
-        )
-
-        client.integrations.create(integration_request)
-        sleep(5)
-        return client.integrations.get_by_alias(integration_alias)
-
-    @staticmethod
-    def cleanup_datasource(client: CodeMieClient, datasource_id: str):
-        """Clean up datasource and verify deletion."""
-        try:
-            client.datasources.delete(datasource_id)
-            sleep(5)
-            with pytest.raises(Exception) as exc_info:
-                client.datasources.get(datasource_id)
-            assert "503" in str(exc_info.value)
-        except Exception as e:
-            pytest.fail(f"Failed to clean up datasource: {str(e)}")
-
     @staticmethod
     def verify_datasource_exists(
-        client: CodeMieClient,
+        search_utils,
         name: str,
         project_name: str,
         datasource_type: DataSourceType,
     ):
         """Verify datasource exists with given parameters."""
-        sleep(5)
-        datasources = client.datasources.list(
+        datasources = search_utils.list_data_sources(
             datasource_types=datasource_type, projects=project_name
         )
         datasource = next((ds for ds in datasources if ds.name == name), None)
-        assert datasource is not None
-        assert datasource.project_name == project_name
-        assert datasource.type == datasource_type
+        assert_that(datasource, is_not(None))
+        assert_that(datasource.project_name, equal_to(project_name))
+        assert_that(datasource.type, equal_to(datasource_type))
         return datasource
 
     @staticmethod
     def verify_datasource_updated(
-        client: CodeMieClient, datasource_id: str, expected_values: dict
+        datasource_utils, datasource_id: str, expected_values: dict
     ):
         """
         Verify datasource was updated with expected values.
         Handles both root level fields and nested objects (Jira, Confluence).
         """
-        sleep(5)
-        updated_datasource = client.datasources.get(datasource_id)
-        assert updated_datasource is not None
+        updated_datasource = datasource_utils.get_datasource(datasource_id)
+        assert_that(updated_datasource, is_not(None))
 
         field_mapping = {
             "jql": ("jira", Jira, "jql"),
@@ -148,15 +95,19 @@ class TestDatasourceBase:
                 nested_obj = getattr(updated_datasource, attr, None)
                 if nested_obj is not None and isinstance(nested_obj, expected_class):
                     actual_value = getattr(nested_obj, sub_attr, None)
-                    assert actual_value == value, (
-                        f"Expected {key} to be {value}, got {actual_value}"
+                    assert_that(
+                        actual_value,
+                        equal_to(value),
+                        f"Expected {key} to be {value}, got {actual_value}",
                     )
                 else:
                     pytest.fail(f"Unhandled field in verification: {key}")
             else:
                 actual_value = getattr(updated_datasource, key, None)
-                assert actual_value == value, (
-                    f"Expected {key} to be {value}, got {actual_value}"
+                assert_that(
+                    actual_value,
+                    equal_to(value),
+                    f"Expected {key} to be {value}, got {actual_value}",
                 )
 
         return updated_datasource
@@ -165,13 +116,16 @@ class TestDatasourceBase:
 class TestDatasources(TestDatasourceBase):
     """Tests for datasource operations."""
 
-    def test_list_datasources(self, client: CodeMieClient):
-        """Test successful retrieval of available datasources models."""
+    def test_list_datasources(self, search_utils):
         datasource_types = [DataSourceType.CODE, DataSourceType.CONFLUENCE]
-        models = client.datasources.list(datasource_types=datasource_types)
-        assert isinstance(models, list)
-        assert len(models) > 0
-        assert all(model.type in datasource_types for model in models)
+        datasource_models = search_utils.list_data_sources(
+            datasource_types=datasource_types
+        )
+        assert_that(
+            datasource_models, all_of(instance_of(list), has_length(greater_than(0)))
+        )
+        for model in datasource_models:
+            assert_that(model.type, is_in(datasource_types))
 
     @pytest.mark.parametrize(
         "datasource_type",
@@ -183,12 +137,15 @@ class TestDatasources(TestDatasourceBase):
             DataSourceType.GOOGLE,
         ],
     )
-    def test_list_datasources_by_type(self, client: CodeMieClient, datasource_type):
-        """Test successful retrieval of available datasources models."""
-        models = client.datasources.list(datasource_types=datasource_type)
-        assert isinstance(models, list)
-        assert len(models) > 0
-        assert all(model.type == datasource_type for model in models)
+    def test_list_datasources_filtered_by_type(self, search_utils, datasource_type):
+        datasource_models = search_utils.list_data_sources(
+            datasource_types=datasource_type
+        )
+        assert_that(
+            datasource_models, all_of(instance_of(list), has_length(greater_than(0)))
+        )
+        for model in datasource_models:
+            assert_that(model.type, is_in(datasource_type))
 
     @pytest.mark.parametrize(
         "status",
@@ -198,194 +155,179 @@ class TestDatasources(TestDatasourceBase):
             DataSourceStatus.FAILED,
         ],
     )
-    def test_list_datasources_by_status(self, client: CodeMieClient, status):
-        """Test successful retrieval of available datasources models."""
-        models = client.datasources.list(status=status)
-        assert isinstance(models, list)
-        assert len(models) > 0
-        assert all(model.status == status for model in models)
+    def test_list_datasources_filtered_by_status(self, search_utils, status):
+        datasource_models = search_utils.list_data_sources(status=status)
+        assert_that(datasource_models, all_of(instance_of(list)))
+        for model in datasource_models:
+            assert_that(
+                model.status,
+                equal_to(status),
+                f"Datasource {model.name} has status {model.status}, expected {status}",
+            )
 
     def test_create_update_code_datasource(
-        self, client: CodeMieClient, integration_config, datasource_name
+        self,
+        integration_utils,
+        llm_utils,
+        search_utils,
+        integration_config,
+        datasource_utils,
     ):
-        project_name = PROJECT
-        integration = self.create_integration(
-            client, project_name, CredentialTypes.GIT, integration_config
+        integration = integration_utils.create_integration(
+            credential_type=CredentialTypes.GIT,
+            credential_values=[
+                CredentialValues(key=k, value=v)
+                for k, v in integration_config[CredentialTypes.GIT].items()
+            ],
+            project_name=PROJECT,
         )
-        embeddings_models = client.llms.list_embeddings()
-        assert len(embeddings_models) > 0
+        embeddings_models = llm_utils.list_embedding_llm_models()
+        assert_that(embeddings_models, has_length(greater_than(0)))
         embeddings_model = embeddings_models[0]
-        datasource = None
-        try:
-            create_request_params = {
-                "name": datasource_name,
-                "project_name": project_name,
-                "description": "Code datasource description",
-                "link": os.getenv("GITLAB_PROJECT"),
-                "branch": "main",
-                "index_type": CodeDataSourceType.CODE,
-                "embeddings_model": embeddings_model.base_name,
-                "setting_id": integration.id,
-            }
-            create_datasource_request = CodeDataSourceRequest(**create_request_params)
-            created = client.datasources.create(create_datasource_request)
-            assert created is not None
-            datasource = self.verify_datasource_exists(
-                client, datasource_name, project_name, DataSourceType.CODE
-            )
+        create_request_params = {
+            "name": integration.alias,
+            "project_name": PROJECT,
+            "description": "Code datasource description",
+            "link": os.getenv("GITLAB_PROJECT"),
+            "branch": "main",
+            "index_type": DataSourceType.CODE,
+            "embeddings_model": embeddings_model.base_name,
+            "setting_id": integration.id,
+        }
+        created = datasource_utils.create_code_datasource(**create_request_params)
+        assert_that(created, is_not(None))
+        datasource = self.verify_datasource_exists(
+            search_utils, integration.alias, PROJECT, DataSourceType.CODE
+        )
 
-            update_request_params = {
-                "link": os.getenv("GITHUB_PROJECT"),
-                "branch": "master",
-                "name": datasource_name,
-                "project_name": project_name,
-                "description": "Updated datasource description",
-            }
+        update_request_params = {
+            "link": os.getenv("GITHUB_PROJECT"),
+            "branch": "master",
+            "name": integration.alias,
+            "project_name": PROJECT,
+            "description": "Updated datasource description",
+        }
 
-            update_datasource_request = UpdateCodeDataSourceRequest(
-                **update_request_params
-            )
-            updated = client.datasources.update(
-                datasource.id, update_datasource_request
-            )
-            assert updated is not None
-            self.verify_datasource_updated(client, datasource.id, update_request_params)
-        finally:
-            # Cleanup
-            if datasource:
-                self.cleanup_datasource(client, datasource.id)
-            if integration:
-                client.integrations.delete(integration.id)
+        updated = datasource_utils.update_code_datasource(
+            datasource.id, **update_request_params
+        )
+        assert_that(updated, is_not(None))
+        self.verify_datasource_updated(
+            datasource_utils, datasource.id, update_request_params
+        )
 
     def test_create_update_confluence_datasource(
-        self, client: CodeMieClient, integration_config, datasource_name
+        self,
+        search_utils,
+        integration_utils,
+        integration_config,
+        datasource_utils,
     ):
-        project_name = PROJECT
-        integration = self.create_integration(
-            client, project_name, CredentialTypes.CONFLUENCE, integration_config
+        integration = integration_utils.create_user_integration(
+            credential_type=CredentialTypes.CONFLUENCE,
+            credentials=[
+                CredentialValues(key=k, value=v)
+                for k, v in integration_config[CredentialTypes.CONFLUENCE].items()
+            ],
+            project_name=PROJECT,
+        )
+        create_request_params = {
+            "name": integration.alias,
+            "project_name": PROJECT,
+            "description": "Datasource for KB space",
+            "cql": os.getenv("CQL"),
+            "setting_id": integration.id,
+        }
+        created = datasource_utils.create_confluence_datasource(**create_request_params)
+        assert_that(created, is_not(None))
+        datasource = self.verify_datasource_exists(
+            search_utils, integration.alias, PROJECT, DataSourceType.CONFLUENCE
         )
-        datasource = None
-        try:
-            create_request_params = {
-                "name": datasource_name,
-                "project_name": project_name,
-                "description": "Datasource for KB space",
-                "cql": os.getenv("CQL"),
-                "setting_id": integration.id,
-            }
-            create_datasource_request = ConfluenceDataSourceRequest(
-                **create_request_params
-            )
-            created = client.datasources.create(create_datasource_request)
-            assert created is not None
-            datasource = self.verify_datasource_exists(
-                client, datasource_name, project_name, DataSourceType.CONFLUENCE
-            )
 
-            update_request_params = {
-                "name": datasource_name,
-                "project_name": project_name,
-                "description": "Updated datasource description for KB space",
-                "cql": "SPACE = MY_KB",
-            }
+        update_request_params = {
+            "name": integration.alias,
+            "description": "Updated datasource description for KB space",
+            "cql": "SPACE = MY_KB",
+        }
 
-            update_datasource_request = UpdateConfluenceDataSourceRequest(
-                **update_request_params
-            )
-            updated = client.datasources.update(
-                datasource.id, update_datasource_request
-            )
-            assert updated is not None
-            self.verify_datasource_updated(client, datasource.id, update_request_params)
-        finally:
-            # Cleanup
-            if datasource:
-                self.cleanup_datasource(client, datasource.id)
-            if integration:
-                client.integrations.delete(integration.id)
+        updated = datasource_utils.update_confluence_datasource(
+            datasource.id, **update_request_params
+        )
+        assert_that(updated, is_not(None))
+        self.verify_datasource_updated(
+            datasource_utils, datasource.id, update_request_params
+        )
 
     def test_create_update_jira_datasource(
-        self, client: CodeMieClient, integration_config, datasource_name
+        self,
+        search_utils,
+        integration_utils,
+        integration_config,
+        datasource_utils,
     ):
-        project_name = PROJECT
-        integration = self.create_integration(
-            client, project_name, CredentialTypes.JIRA, integration_config
+        integration = integration_utils.create_integration(
+            credential_type=CredentialTypes.JIRA,
+            credential_values=[
+                CredentialValues(key=k, value=v)
+                for k, v in integration_config[CredentialTypes.JIRA].items()
+            ],
+            setting_type=IntegrationType.USER,
+            project_name=PROJECT,
+        )
+        create_request_params = {
+            "name": integration.alias,
+            "project_name": PROJECT,
+            "description": "Jira datasource description",
+            "jql": os.getenv("JQL"),
+            "setting_id": integration.id,
+        }
+        created = datasource_utils.create_jira_datasource(**create_request_params)
+        assert_that(created, is_not(None))
+        datasource = self.verify_datasource_exists(
+            search_utils, integration.alias, PROJECT, DataSourceType.JIRA
         )
-        datasource = None
-        try:
-            create_request_params = {
-                "name": datasource_name,
-                "project_name": project_name,
-                "description": "Jira datasource description",
-                "jql": os.getenv("JQL"),
-                "setting_id": integration.id,
-            }
-            create_datasource_request = JiraDataSourceRequest(**create_request_params)
-            created = client.datasources.create(create_datasource_request)
-            assert created is not None
-            datasource = self.verify_datasource_exists(
-                client, datasource_name, project_name, DataSourceType.JIRA
-            )
 
-            update_request_params = {
-                "name": datasource_name,
-                "project_name": project_name,
-                "description": "Updated Jira datasource description",
-                "jql": os.getenv("JQL"),
-            }
+        update_request_params = {
+            "name": integration.alias,
+            "project_name": PROJECT,
+            "description": "Updated Jira datasource description",
+            "jql": os.getenv("JQL"),
+        }
 
-            update_datasource_request = UpdateJiraDataSourceRequest(
-                **update_request_params
-            )
-            updated = client.datasources.update(
-                datasource.id, update_datasource_request
-            )
-            assert updated is not None
-            self.verify_datasource_updated(client, datasource.id, update_request_params)
-        finally:
-            # Cleanup
-            if datasource:
-                self.cleanup_datasource(client, datasource.id)
-            if integration:
-                client.integrations.delete(integration.id)
+        updated = datasource_utils.update_jira_datasource(
+            datasource.id, **update_request_params
+        )
+        assert_that(updated, is_not(None))
+        self.verify_datasource_updated(
+            datasource_utils, datasource.id, update_request_params
+        )
 
-    def test_create_update_google_datasource(
-        self, client: CodeMieClient, integration_config, datasource_name
-    ):
-        project_name = PROJECT
-        datasource = None
-        try:
-            create_request_params = {
-                "name": datasource_name,
-                "project_name": project_name,
-                "description": "Google datasource description",
-                "google_doc": "https://docs.google.com/document/d/16qP3UlOKa-wFA2Ztbsu6qfO4jGPM6S65bse5C9D-AzQ/edit?tab=t.0",
-            }
-            create_datasource_request = GoogleDataSourceRequest(**create_request_params)
-            created = client.datasources.create(create_datasource_request)
-            assert created is not None
-            datasource = self.verify_datasource_exists(
-                client, datasource_name, project_name, DataSourceType.GOOGLE
-            )
+    def test_create_update_google_datasource(self, search_utils, datasource_utils):
+        create_request_params = {
+            "name": get_random_name(),
+            "project_name": PROJECT,
+            "description": "Google datasource description",
+            "google_doc": "https://docs.google.com/document/d/19EXgnFCgJontz0ToCAH6zMGwBTdhi5X97P9JIby4wHs/edit?tab=t.0",
+        }
+        created = datasource_utils.create_google_doc_datasource(**create_request_params)
+        assert_that(created, is_not(None))
+        datasource = self.verify_datasource_exists(
+            search_utils, created.name, PROJECT, DataSourceType.GOOGLE
+        )
 
-            update_request_params = {
-                "name": datasource_name,
-                "project_name": project_name,
-                "description": "Updated Google datasource description",
-            }
+        update_request_params = {
+            "name": created.name,
+            "project_name": PROJECT,
+            "description": "Updated Google datasource description",
+        }
 
-            update_datasource_request = UpdateGoogleDataSourceRequest(
-                **update_request_params
-            )
-            updated = client.datasources.update(
-                datasource.id, update_datasource_request
-            )
-            assert updated is not None
-            self.verify_datasource_updated(client, datasource.id, update_request_params)
-        finally:
-            # Cleanup
-            if datasource:
-                self.cleanup_datasource(client, datasource.id)
+        updated = datasource_utils.update_google_doc_datasource(
+            datasource.id, **update_request_params
+        )
+        assert_that(updated, is_not(None))
+        self.verify_datasource_updated(
+            datasource_utils, datasource.id, update_request_params
+        )
 
     @pytest.mark.parametrize(
         "datasource_type",
@@ -397,54 +339,96 @@ class TestDatasources(TestDatasourceBase):
             DataSourceType.GOOGLE,
         ],
     )
-    def test_get_datasource_by_id(self, client: CodeMieClient, datasource_type):
-        """Test successful retrieval of available datasources models."""
-        datasources = client.datasources.list(
+    def test_get_datasource_by_id(
+        self, datasource_utils, datasource_type, search_utils
+    ):
+        datasources = search_utils.list_data_sources(
             datasource_types=datasource_type, per_page=50
         )
-        assert isinstance(datasources, list)
-        assert len(datasources) > 0
+        assert_that(datasources, all_of(instance_of(list), has_length(greater_than(0))))
 
         original_datasource = datasources[0]
         datasource_id = original_datasource.id
-        retrieved_datasource = client.datasources.get(datasource_id)
+        retrieved_datasource = datasource_utils.get_datasource(datasource_id)
 
         # Compare full objects (they should be identical)
-        assert retrieved_datasource.id == original_datasource.id
-        assert retrieved_datasource.name == original_datasource.name
-        assert retrieved_datasource.project_name == original_datasource.project_name
-        assert retrieved_datasource.created_by == original_datasource.created_by
-        assert (
-            retrieved_datasource.shared_with_project
-            == original_datasource.shared_with_project
+        assert_that(
+            retrieved_datasource,
+            all_of(
+                is_not(None),
+            ),
+        )
+        assert_that(retrieved_datasource.id, equal_to(original_datasource.id))
+        assert_that(retrieved_datasource.name, equal_to(original_datasource.name))
+        assert_that(
+            retrieved_datasource.project_name,
+            equal_to(original_datasource.project_name),
+        )
+        assert_that(
+            retrieved_datasource.created_by,
+            equal_to(original_datasource.created_by),
         )
-        assert retrieved_datasource.created_date == original_datasource.created_date
-        assert retrieved_datasource.error_message == original_datasource.error_message
-        assert retrieved_datasource.processing_info is not None
-        assert (
-            retrieved_datasource.processing_info.processed_documents_count is not None
+        assert_that(
+            retrieved_datasource.shared_with_project,
+            equal_to(original_datasource.shared_with_project),
+        )
+        assert_that(
+            retrieved_datasource.created_date,
+            equal_to(original_datasource.created_date),
+        )
+        assert_that(
+            retrieved_datasource.error_message,
+            equal_to(original_datasource.error_message),
+        )
+        assert_that(
+            retrieved_datasource.processing_info.processed_documents_count,
+            is_not(None),
         )
 
         if datasource_type == DataSourceType.CODE:
-            assert original_datasource.description is None
-            assert retrieved_datasource.confluence is None
-            assert retrieved_datasource.jira is None
-            assert retrieved_datasource.tokens_usage is not None
-            assert retrieved_datasource.code.link == original_datasource.code.link
-            assert retrieved_datasource.code.branch is not None
+            assert_that(original_datasource.description, is_(None))
+            assert_that(
+                retrieved_datasource,
+                all_of(
+                    has_property("confluence", None),
+                    has_property("jira", None),
+                    has_property("tokens_usage", is_not(None)),
+                    has_property("code"),
+                ),
+            )
+            assert_that(
+                retrieved_datasource.code.link,
+                equal_to(original_datasource.code.link),
+            )
+            assert_that(retrieved_datasource.code.branch, is_not(None))
         elif datasource_type == DataSourceType.CONFLUENCE:
-            assert retrieved_datasource.code is None
-            assert retrieved_datasource.jira is None
-            assert retrieved_datasource.tokens_usage is not None
-            assert retrieved_datasource.confluence is not None
-            assert retrieved_datasource.confluence.cql is not None
+            assert_that(
+                retrieved_datasource,
+                all_of(
+                    has_property("code", None),
+                    has_property("jira", None),
+                    has_property("tokens_usage", is_not(None)),
+                    has_property("confluence", is_not(None)),
+                ),
+            )
+            assert_that(retrieved_datasource.confluence.cql, is_not(None))
         elif datasource_type == DataSourceType.JIRA:
-            assert retrieved_datasource.code is None
-            assert retrieved_datasource.confluence is None
-            assert retrieved_datasource.jira is not None
-            assert retrieved_datasource.jira.jql is not None
+            assert_that(
+                retrieved_datasource,
+                all_of(
+                    has_property("code", None),
+                    has_property("confluence", None),
+                    has_property("jira", is_not(None)),
+                ),
+            )
+            assert_that(retrieved_datasource.jira.jql, is_not(None))
         elif datasource_type == DataSourceType.GOOGLE:
-            assert retrieved_datasource.code is None
-            assert retrieved_datasource.confluence is None
-            assert retrieved_datasource.jira is None
-            assert retrieved_datasource.google_doc_link is not None
+            assert_that(
+                retrieved_datasource,
+                all_of(
+                    has_property("code", None),
+                    has_property("confluence", None),
+                    has_property("jira", None),
+                    has_property("google_doc_link", is_not(None)),
+                ),
+            )