passing unit tests

databricks · Feb 14, 2025 · 722c758 · 722c758
1 parent 5086064
commit 722c758
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 16 deletions.
diff --git a/dbt/adapters/databricks/connections.py b/dbt/adapters/databricks/connections.py
@@ -25,22 +25,20 @@
 )
 from dbt.adapters.databricks.__version__ import version as __version__
 from dbt.adapters.databricks.api_client import DatabricksApiClient
-from dbt.adapters.databricks.connections import (
-    ConnectionCreate,
-    ConnectionCreateError,
-    ConnectionIdleCheck,
-    ConnectionIdleClose,
-    ConnectionRetrieve,
-    ConnectionReuse,
-)
 from dbt.adapters.databricks.credentials import (
     DatabricksCredentialManager,
     DatabricksCredentials,
 )
 from dbt.adapters.databricks.events.connection_events import (
     ConnectionAcquire,
+    ConnectionCreate,
+    ConnectionCreateError,
+    ConnectionIdleCheck,
+    ConnectionIdleClose,
     ConnectionRelease,
     ConnectionReset,
+    ConnectionRetrieve,
+    ConnectionReuse,
 )
 from dbt.adapters.databricks.events.other_events import QueryError
 from dbt.adapters.databricks.handle import CursorWrapper, DatabricksHandle, SqlUtils

diff --git a/dbt/adapters/databricks/handle.py b/dbt/adapters/databricks/handle.py
@@ -12,7 +12,7 @@
 from dbt.adapters.contracts.connection import AdapterResponse
 from dbt.adapters.databricks import utils
 from dbt.adapters.databricks.__version__ import version as __version__
-from dbt.adapters.databricks.credentials import DatabricksCredentials, TCredentialProvider
+from dbt.adapters.databricks.credentials import DatabricksCredentialManager, DatabricksCredentials
 from dbt.adapters.databricks.logging import logger
 
 if TYPE_CHECKING:
@@ -297,7 +297,7 @@ def clean_sql(sql: str) -> str:
 
     @staticmethod
     def prepare_connection_arguments(
-        creds: DatabricksCredentials, creds_provider: TCredentialProvider, http_path: str
+        creds: DatabricksCredentials, creds_manager: DatabricksCredentialManager, http_path: str
     ) -> dict[str, Any]:
         invocation_env = creds.get_invocation_env()
         user_agent_entry = SqlUtils.user_agent
@@ -313,7 +313,7 @@ def prepare_connection_arguments(
         return {
             "server_hostname": creds.host,
             "http_path": http_path,
-            "credentials_provider": creds_provider,
+            "credentials_provider": creds_manager.credentials_provider,
             "http_headers": http_headers if http_headers else None,
             "session_configuration": creds.session_properties,
             "catalog": creds.database,

diff --git a/tests/unit/test_constraints.py b/tests/unit/test_constraints.py
@@ -24,7 +24,7 @@
 )
 
 
-class TestConstraint(TypedConstraint):
+class FakeConstraint(TypedConstraint):
     str_type = "unique"
 
     def _render_suffix(self):
@@ -34,10 +34,10 @@ def _render_suffix(self):
 class TestTypedConstraint:
     def test_typed_constraint_from_dict__invalid_type(self):
         with pytest.raises(AssertionError, match="Mismatched constraint type"):
-            TestConstraint.from_dict({"type": "custom"})
+            FakeConstraint.from_dict({"type": "custom"})
 
     def test_render__with_name(self):
-        constraint = TestConstraint(type=ConstraintType.check, name="my_constraint")
+        constraint = FakeConstraint(type=ConstraintType.check, name="my_constraint")
         assert constraint.render() == "CONSTRAINT my_constraint test"
 
 

diff --git a/tests/unit/test_query_config.py b/tests/unit/test_query_config.py
@@ -1,9 +1,10 @@
+from unittest.mock import Mock
+
 import pytest
 from dbt_common.exceptions import DbtRuntimeError
 
 from dbt.adapters.databricks import connections
 from dbt.adapters.databricks.connections import QueryConfigUtils, QueryContextWrapper
-from dbt.adapters.databricks.credentials import DatabricksCredentials
 
 
 class TestQueryConfigUtils:
@@ -22,7 +23,7 @@ def path(self):
 
     @pytest.fixture
     def creds(self, path):
-        return DatabricksCredentials(http_path=path)
+        return Mock(http_path=path, compute={}, connect_max_idle=None)
 
     def test_get_http_path__empty(self, path, creds):
         assert QueryConfigUtils.get_http_path(QueryContextWrapper(), creds) == path