[tests] Fixed unit tests and enabled CI job (qa-checks, unit-tests)

Aryamanz29 · Aryamanz29 · commit 493ae5297076 · 2025-06-03T16:50:55.000+05:30
diff --git a/.github/workflows/pyatlan-pr.yaml b/.github/workflows/pyatlan-pr.yaml
@@ -60,16 +60,16 @@ jobs:
           if [ -f requirements.txt ]; then pip install --no-cache-dir -r requirements.txt; fi
           if [ -f requirements-dev.txt ]; then pip install --no-cache-dir -r requirements-dev.txt; fi
 
-      # - name: QA checks (ruff-format, ruff-lint, mypy)
-      #   run: |
-      #     ./qa-checks
-
-      # - name: Run unit tests
-      #   env: # Test tenant environment variables
-      #     ATLAN_API_KEY: ${{ secrets.ATLAN_API_KEY }}
-      #     ATLAN_BASE_URL: ${{ secrets.ATLAN_BASE_URL }}
-      #   # Run with `pytest-sugar` for enhancing the overall test report output
-      #   run: pytest tests/unit --force-sugar
+      - name: QA checks (ruff-format, ruff-lint, mypy)
+        run: |
+          ./qa-checks
+
+      - name: Run unit tests
+        env: # Test tenant environment variables
+          ATLAN_API_KEY: ${{ secrets.ATLAN_API_KEY }}
+          ATLAN_BASE_URL: ${{ secrets.ATLAN_BASE_URL }}
+        # Run with `pytest-sugar` for enhancing the overall test report output
+        run: pytest tests/unit --force-sugar
 
       - name: Prepare integration tests distribution
         id: distribute-integration-test-files
diff --git a/pyatlan/model/packages/base/crawler.py b/pyatlan/model/packages/base/crawler.py
@@ -2,6 +2,7 @@
 from typing import Any, Dict, List, Optional
 
 from pyatlan import utils
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.errors import ErrorCode
 from pyatlan.model.assets import Connection
 from pyatlan.model.packages.base.package import AbstractPackage
@@ -27,6 +28,7 @@ class AbstractCrawler(AbstractPackage):
 
     def __init__(
         self,
+        client: AtlanClient,
         connection_name: str,
         connection_type: str,
         admin_roles: Optional[List[str]] = None,
@@ -38,6 +40,7 @@ def __init__(
         source_logo: str = "",
     ):
         super().__init__()
+        self._client = client
         self._connection_name = connection_name
         self._connection_type = connection_type
         self._admin_roles = admin_roles
@@ -55,6 +58,7 @@ def _get_connection(self) -> Connection:
         which will be the target for the package to crawl assets.
         """
         connection = Connection.create(
+            client=self._client,
             name=self._connection_name,
             connector_type=self._connection_type,
             admin_roles=self._admin_roles,
diff --git a/pyatlan/model/packages/big_query_crawler.py b/pyatlan/model/packages/big_query_crawler.py
@@ -2,6 +2,7 @@
 
 from typing import Dict, List, Optional
 
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.enums import AtlanConnectorType, WorkflowPackage
 from pyatlan.model.packages.base.crawler import AbstractCrawler
 from pyatlan.model.workflow import WorkflowMetadata
@@ -32,6 +33,7 @@ class BigQueryCrawler(AbstractCrawler):
 
     def __init__(
         self,
+        client: AtlanClient,
         connection_name: str,
         admin_roles: Optional[List[str]] = None,
         admin_groups: Optional[List[str]] = None,
@@ -42,6 +44,7 @@ def __init__(
     ):
         self._advanced_config = False
         super().__init__(
+            client=client,
             connection_name=connection_name,
             connection_type=self._CONNECTOR_TYPE,
             admin_roles=admin_roles,
diff --git a/pyatlan/model/packages/confluent_kafka_crawler.py b/pyatlan/model/packages/confluent_kafka_crawler.py
@@ -2,6 +2,7 @@
 
 from typing import List, Optional
 
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.enums import AtlanConnectorType, WorkflowPackage
 from pyatlan.model.packages.base.crawler import AbstractCrawler
 from pyatlan.model.workflow import WorkflowMetadata
@@ -32,6 +33,7 @@ class ConfluentKafkaCrawler(AbstractCrawler):
 
     def __init__(
         self,
+        client: AtlanClient,
         connection_name: str,
         admin_roles: Optional[List[str]] = None,
         admin_groups: Optional[List[str]] = None,
@@ -41,6 +43,7 @@ def __init__(
         row_limit: int = 0,
     ):
         super().__init__(
+            client=client,
             connection_name=connection_name,
             connection_type=self._CONNECTOR_TYPE,
             admin_roles=admin_roles,
diff --git a/pyatlan/model/packages/databricks_crawler.py b/pyatlan/model/packages/databricks_crawler.py
@@ -3,6 +3,7 @@
 from enum import Enum
 from typing import List, Optional
 
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.enums import AtlanConnectorType, WorkflowPackage
 from pyatlan.model.packages.base.crawler import AbstractCrawler
 from pyatlan.model.workflow import WorkflowMetadata
@@ -37,6 +38,7 @@ class ExtractionMethod(str, Enum):
 
     def __init__(
         self,
+        client: AtlanClient,
         connection_name: str,
         admin_roles: Optional[List[str]] = None,
         admin_groups: Optional[List[str]] = None,
@@ -47,6 +49,7 @@ def __init__(
     ):
         self._advanced_config = False
         super().__init__(
+            client=client,
             connection_name=connection_name,
             connection_type=self._CONNECTOR_TYPE,
             admin_roles=admin_roles,
diff --git a/pyatlan/model/packages/dbt_crawler.py b/pyatlan/model/packages/dbt_crawler.py
@@ -2,6 +2,7 @@
 
 from typing import List, Optional
 
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.enums import AtlanConnectorType, WorkflowPackage
 from pyatlan.model.packages.base.crawler import AbstractCrawler
 from pyatlan.model.workflow import WorkflowMetadata
@@ -32,6 +33,7 @@ class DbtCrawler(AbstractCrawler):
 
     def __init__(
         self,
+        client: AtlanClient,
         connection_name: str,
         admin_roles: Optional[List[str]] = None,
         admin_groups: Optional[List[str]] = None,
@@ -41,6 +43,7 @@ def __init__(
         row_limit: int = 0,
     ):
         super().__init__(
+            client=client,
             connection_name=connection_name,
             connection_type=self._CONNECTOR_TYPE,
             admin_roles=admin_roles,
diff --git a/pyatlan/model/packages/dynamo_d_b_crawler.py b/pyatlan/model/packages/dynamo_d_b_crawler.py
@@ -2,6 +2,7 @@
 
 from typing import List, Optional
 
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.enums import AtlanConnectorType, WorkflowPackage
 from pyatlan.model.packages.base.crawler import AbstractCrawler
 from pyatlan.model.workflow import WorkflowMetadata
@@ -32,6 +33,7 @@ class DynamoDBCrawler(AbstractCrawler):
 
     def __init__(
         self,
+        client: AtlanClient,
         connection_name: str,
         admin_roles: Optional[List[str]] = None,
         admin_groups: Optional[List[str]] = None,
@@ -41,6 +43,7 @@ def __init__(
         row_limit: int = 10000,
     ):
         super().__init__(
+            client=client,
             connection_name=connection_name,
             connection_type=self._CONNECTOR_TYPE,
             admin_roles=admin_roles,
diff --git a/pyatlan/model/packages/glue_crawler.py b/pyatlan/model/packages/glue_crawler.py
@@ -3,6 +3,7 @@
 from json import dumps
 from typing import List, Optional
 
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.errors import ErrorCode
 from pyatlan.model.enums import AtlanConnectorType, WorkflowPackage
 from pyatlan.model.packages.base.crawler import AbstractCrawler
@@ -39,6 +40,7 @@ class GlueCrawler(AbstractCrawler):
 
     def __init__(
         self,
+        client: AtlanClient,
         connection_name: str,
         admin_roles: Optional[List[str]] = None,
         admin_groups: Optional[List[str]] = None,
@@ -48,6 +50,7 @@ def __init__(
         row_limit: int = 0,
     ):
         super().__init__(
+            client=client,
             connection_name=connection_name,
             connection_type=self._CONNECTOR_TYPE,
             admin_roles=admin_roles,
diff --git a/pyatlan/model/packages/mongodb_crawler.py b/pyatlan/model/packages/mongodb_crawler.py
@@ -2,6 +2,7 @@
 
 from typing import Dict, List, Optional
 
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.enums import AtlanConnectorType, WorkflowPackage
 from pyatlan.model.packages.base.crawler import AbstractCrawler
 from pyatlan.model.workflow import WorkflowMetadata
@@ -32,6 +33,7 @@ class MongoDBCrawler(AbstractCrawler):
 
     def __init__(
         self,
+        client: AtlanClient,
         connection_name: str,
         admin_roles: Optional[List[str]] = None,
         admin_groups: Optional[List[str]] = None,
@@ -41,6 +43,7 @@ def __init__(
         row_limit: int = 10000,
     ):
         super().__init__(
+            client=client,
             connection_name=connection_name,
             connection_type=self._CONNECTOR_TYPE,
             admin_roles=admin_roles,
diff --git a/pyatlan/model/packages/oracle_crawler.py b/pyatlan/model/packages/oracle_crawler.py
@@ -4,6 +4,7 @@
 from json import dumps
 from typing import Any, Dict, List, Optional
 
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.enums import AtlanConnectorType, WorkflowPackage
 from pyatlan.model.packages.base.crawler import AbstractCrawler
 from pyatlan.model.workflow import WorkflowMetadata
@@ -58,6 +59,7 @@ class SecretStore(str, Enum):
 
     def __init__(
         self,
+        client: AtlanClient,
         connection_name: str,
         admin_roles: Optional[List[str]] = None,
         admin_groups: Optional[List[str]] = None,
@@ -69,6 +71,7 @@ def __init__(
         self._advanced_config = False
         self._agent_config = False
         super().__init__(
+            client=client,
             connection_name=connection_name,
             connection_type=self._CONNECTOR_TYPE,
             admin_roles=admin_roles,
diff --git a/pyatlan/model/packages/postgres_crawler.py b/pyatlan/model/packages/postgres_crawler.py
@@ -2,6 +2,7 @@
 
 from typing import List, Optional
 
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.enums import AtlanConnectorType, WorkflowPackage
 from pyatlan.model.packages.base.crawler import AbstractCrawler
 from pyatlan.model.workflow import WorkflowMetadata
@@ -32,6 +33,7 @@ class PostgresCrawler(AbstractCrawler):
 
     def __init__(
         self,
+        client: AtlanClient,
         connection_name: str,
         admin_roles: Optional[List[str]] = None,
         admin_groups: Optional[List[str]] = None,
@@ -41,6 +43,7 @@ def __init__(
         row_limit: int = 10000,
     ):
         super().__init__(
+            client=client,
             connection_name=connection_name,
             connection_type=self._CONNECTOR_TYPE,
             admin_roles=admin_roles,
diff --git a/pyatlan/model/packages/powerbi_crawler.py b/pyatlan/model/packages/powerbi_crawler.py
@@ -2,6 +2,7 @@
 
 from typing import List, Optional
 
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.enums import AtlanConnectorType, WorkflowPackage
 from pyatlan.model.packages.base.crawler import AbstractCrawler
 from pyatlan.model.workflow import WorkflowMetadata
@@ -36,6 +37,7 @@ class PowerBICrawler(AbstractCrawler):
 
     def __init__(
         self,
+        client: AtlanClient,
         connection_name: str,
         admin_roles: Optional[List[str]] = None,
         admin_groups: Optional[List[str]] = None,
@@ -45,6 +47,7 @@ def __init__(
         row_limit: int = 0,
     ):
         super().__init__(
+            client=client,
             connection_name=connection_name,
             connection_type=self._CONNECTOR_TYPE,
             admin_roles=admin_roles,
diff --git a/pyatlan/model/packages/s_q_l_server_crawler.py b/pyatlan/model/packages/s_q_l_server_crawler.py
@@ -2,6 +2,7 @@
 
 from typing import List, Optional
 
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.enums import AtlanConnectorType, WorkflowPackage
 from pyatlan.model.packages.base.crawler import AbstractCrawler
 from pyatlan.model.workflow import WorkflowMetadata
@@ -32,6 +33,7 @@ class SQLServerCrawler(AbstractCrawler):
 
     def __init__(
         self,
+        client: AtlanClient,
         connection_name: str,
         admin_roles: Optional[List[str]] = None,
         admin_groups: Optional[List[str]] = None,
@@ -41,6 +43,7 @@ def __init__(
         row_limit: int = 10000,
     ):
         super().__init__(
+            client=client,
             connection_name=connection_name,
             connection_type=self._CONNECTOR_TYPE,
             admin_roles=admin_roles,
diff --git a/pyatlan/model/packages/sigma_crawler.py b/pyatlan/model/packages/sigma_crawler.py
@@ -3,6 +3,7 @@
 from enum import Enum
 from typing import List, Optional
 
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.enums import AtlanConnectorType, WorkflowPackage
 from pyatlan.model.packages.base.crawler import AbstractCrawler
 from pyatlan.model.workflow import WorkflowMetadata
@@ -40,6 +41,7 @@ class Hostname(str, Enum):
 
     def __init__(
         self,
+        client: AtlanClient,
         connection_name: str,
         admin_roles: Optional[List[str]] = None,
         admin_groups: Optional[List[str]] = None,
@@ -49,6 +51,7 @@ def __init__(
         row_limit: int = 0,
     ):
         super().__init__(
+            client=client,
             connection_name=connection_name,
             connection_type=self._CONNECTOR_TYPE,
             admin_roles=admin_roles,
diff --git a/pyatlan/model/packages/snowflake_crawler.py b/pyatlan/model/packages/snowflake_crawler.py
@@ -2,6 +2,7 @@
 
 from typing import List, Optional
 
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.enums import AtlanConnectorType, WorkflowPackage
 from pyatlan.model.packages.base.crawler import AbstractCrawler
 from pyatlan.model.workflow import WorkflowMetadata
@@ -32,6 +33,7 @@ class SnowflakeCrawler(AbstractCrawler):
 
     def __init__(
         self,
+        client: AtlanClient,
         connection_name: str,
         admin_roles: Optional[List[str]] = None,
         admin_groups: Optional[List[str]] = None,
@@ -41,6 +43,7 @@ def __init__(
         row_limit: int = 10000,
     ):
         super().__init__(
+            client=client,
             connection_name=connection_name,
             connection_type=self._CONNECTOR_TYPE,
             admin_roles=admin_roles,
diff --git a/pyatlan/model/packages/sql_server_crawler.py b/pyatlan/model/packages/sql_server_crawler.py
@@ -2,6 +2,7 @@
 
 from typing import List, Optional
 
+from pyatlan.client.atlan import AtlanClient
 from pyatlan.model.enums import AtlanConnectorType, WorkflowPackage
 from pyatlan.model.packages.base.crawler import AbstractCrawler
 from pyatlan.model.workflow import WorkflowMetadata
@@ -32,6 +33,7 @@ class SQLServerCrawler(AbstractCrawler):
 
     def __init__(
         self,
+        client: AtlanClient,
         connection_name: str,
         admin_roles: Optional[List[str]] = None,
         admin_groups: Optional[List[str]] = None,
@@ -41,6 +43,7 @@ def __init__(
         row_limit: int = 10000,
     ):
         super().__init__(
+            client=client,
             connection_name=connection_name,
             connection_type=self._CONNECTOR_TYPE,
             admin_roles=admin_roles,
diff --git a/pyatlan/model/packages/tableau_crawler.py b/pyatlan/model/packages/tableau_crawler.py
diff --git a/tests/unit/model/badge_test.py b/tests/unit/model/badge_test.py
diff --git a/tests/unit/model/connection_test.py b/tests/unit/model/connection_test.py
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
diff --git a/tests/unit/test_custom_metadata.py b/tests/unit/test_custom_metadata.py
diff --git a/tests/unit/test_packages.py b/tests/unit/test_packages.py
diff --git a/tests/unit/test_source_cache.py b/tests/unit/test_source_cache.py
diff --git a/tests/unit/test_typedef_model.py b/tests/unit/test_typedef_model.py
