Added Databricks Connect fixture

databrickslabs · Sep 16, 2024 · 53034fd · 53034fd
1 parent 7fea983
commit 53034fd
Show file tree

Hide file tree

Showing 6 changed files with 65 additions and 6 deletions.
diff --git a/README.md b/README.md
@@ -185,7 +185,7 @@ Loads environment variables specified in [`debug_env_name` fixture](#debug_env_n
 for local debugging in IDEs, otherwise allowing the tests to run with the default environment variables
 specified in the CI/CD pipeline.
 
-See also [`env_or_skip`](#env_or_skip-fixture), [`ws`](#ws-fixture), [`debug_env_name`](#debug_env_name-fixture).
+See also [`env_or_skip`](#env_or_skip-fixture), [`ws`](#ws-fixture), [`debug_env_name`](#debug_env_name-fixture), [`is_in_debug`](#is_in_debug-fixture).
 
 
 [[back to top](#python-testing-for-databricks)]
@@ -204,7 +204,7 @@ def test_something(env_or_skip):
     assert token is not None
 ```
 
-See also [`make_udf`](#make_udf-fixture), [`sql_backend`](#sql_backend-fixture), [`debug_env`](#debug_env-fixture).
+See also [`make_udf`](#make_udf-fixture), [`sql_backend`](#sql_backend-fixture), [`debug_env`](#debug_env-fixture), [`is_in_debug`](#is_in_debug-fixture).
 
 
 [[back to top](#python-testing-for-databricks)]
@@ -227,7 +227,22 @@ def test_workspace_operations(ws):
     assert len(clusters) >= 0
 ```
 
-See also [`log_workspace_link`](#log_workspace_link-fixture), [`make_alert_permissions`](#make_alert_permissions-fixture), [`make_authorization_permissions`](#make_authorization_permissions-fixture), [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster-fixture), [`make_cluster_permissions`](#make_cluster_permissions-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_cluster_policy_permissions`](#make_cluster_policy_permissions-fixture), [`make_dashboard_permissions`](#make_dashboard_permissions-fixture), [`make_directory`](#make_directory-fixture), [`make_directory_permissions`](#make_directory_permissions-fixture), [`make_experiment`](#make_experiment-fixture), [`make_experiment_permissions`](#make_experiment_permissions-fixture), [`make_feature_table_permissions`](#make_feature_table_permissions-fixture), [`make_group`](#make_group-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_instance_pool_permissions`](#make_instance_pool_permissions-fixture), [`make_job`](#make_job-fixture), [`make_job_permissions`](#make_job_permissions-fixture), [`make_lakeview_dashboard_permissions`](#make_lakeview_dashboard_permissions-fixture), [`make_model`](#make_model-fixture), [`make_notebook`](#make_notebook-fixture), [`make_notebook_permissions`](#make_notebook_permissions-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_pipeline_permissions`](#make_pipeline_permissions-fixture), [`make_query`](#make_query-fixture), [`make_query_permissions`](#make_query_permissions-fixture), [`make_registered_model_permissions`](#make_registered_model_permissions-fixture), [`make_repo`](#make_repo-fixture), [`make_repo_permissions`](#make_repo_permissions-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_secret_scope_acl`](#make_secret_scope_acl-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_serving_endpoint_permissions`](#make_serving_endpoint_permissions-fixture), [`make_storage_credential`](#make_storage_credential-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_warehouse`](#make_warehouse-fixture), [`make_warehouse_permissions`](#make_warehouse_permissions-fixture), [`make_workspace_file_path_permissions`](#make_workspace_file_path_permissions-fixture), [`make_workspace_file_permissions`](#make_workspace_file_permissions-fixture), [`sql_backend`](#sql_backend-fixture), [`debug_env`](#debug_env-fixture), [`product_info`](#product_info-fixture).
+See also [`log_workspace_link`](#log_workspace_link-fixture), [`make_alert_permissions`](#make_alert_permissions-fixture), [`make_authorization_permissions`](#make_authorization_permissions-fixture), [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster-fixture), [`make_cluster_permissions`](#make_cluster_permissions-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_cluster_policy_permissions`](#make_cluster_policy_permissions-fixture), [`make_dashboard_permissions`](#make_dashboard_permissions-fixture), [`make_directory`](#make_directory-fixture), [`make_directory_permissions`](#make_directory_permissions-fixture), [`make_experiment`](#make_experiment-fixture), [`make_experiment_permissions`](#make_experiment_permissions-fixture), [`make_feature_table_permissions`](#make_feature_table_permissions-fixture), [`make_group`](#make_group-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_instance_pool_permissions`](#make_instance_pool_permissions-fixture), [`make_job`](#make_job-fixture), [`make_job_permissions`](#make_job_permissions-fixture), [`make_lakeview_dashboard_permissions`](#make_lakeview_dashboard_permissions-fixture), [`make_model`](#make_model-fixture), [`make_notebook`](#make_notebook-fixture), [`make_notebook_permissions`](#make_notebook_permissions-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_pipeline_permissions`](#make_pipeline_permissions-fixture), [`make_query`](#make_query-fixture), [`make_query_permissions`](#make_query_permissions-fixture), [`make_registered_model_permissions`](#make_registered_model_permissions-fixture), [`make_repo`](#make_repo-fixture), [`make_repo_permissions`](#make_repo_permissions-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_secret_scope_acl`](#make_secret_scope_acl-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_serving_endpoint_permissions`](#make_serving_endpoint_permissions-fixture), [`make_storage_credential`](#make_storage_credential-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_warehouse`](#make_warehouse-fixture), [`make_warehouse_permissions`](#make_warehouse_permissions-fixture), [`make_workspace_file_path_permissions`](#make_workspace_file_path_permissions-fixture), [`make_workspace_file_permissions`](#make_workspace_file_permissions-fixture), [`spark`](#spark-fixture), [`sql_backend`](#sql_backend-fixture), [`debug_env`](#debug_env-fixture), [`product_info`](#product_info-fixture).
+
+
+[[back to top](#python-testing-for-databricks)]
+
+### `spark` fixture
+Get Databricks Connect Spark session. Requires `databricks-connect` package to be installed.
+
+Usage:
+```python
+def test_databricks_connect(spark):
+    rows = spark.sql("SELECT 1").collect()
+    assert rows[0][0] == 1
+```
+
+See also [`ws`](#ws-fixture).
 
 
 [[back to top](#python-testing-for-databricks)]
@@ -990,6 +1005,14 @@ HEX-encoded purge time suffix for test objects.
 See also [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_directory`](#make_directory-fixture), [`make_experiment`](#make_experiment-fixture), [`make_group`](#make_group-fixture), [`make_notebook`](#make_notebook-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_query`](#make_query-fixture), [`make_repo`](#make_repo-fixture), [`make_user`](#make_user-fixture), [`watchdog_remove_after`](#watchdog_remove_after-fixture).
 
 
+[[back to top](#python-testing-for-databricks)]
+
+### `is_in_debug` fixture
+_No description yet._
+
+See also [`debug_env`](#debug_env-fixture), [`env_or_skip`](#env_or_skip-fixture).
+
+
 [[back to top](#python-testing-for-databricks)]
 
 <!-- END FIXTURES -->

diff --git a/src/databricks/labs/pytester/fixtures/connect.py b/src/databricks/labs/pytester/fixtures/connect.py
@@ -0,0 +1,33 @@
+import logging
+
+from pytest import fixture, skip
+from databricks.sdk import WorkspaceClient
+
+logger = logging.getLogger(__name__)
+
+
+@fixture
+def spark(ws: WorkspaceClient):
+    """
+    Get Databricks Connect Spark session. Requires `databricks-connect` package to be installed.
+
+    Usage:
+    ```python
+    def test_databricks_connect(spark):
+        rows = spark.sql("SELECT 1").collect()
+        assert rows[0][0] == 1
+    ```
+    """
+    if not ws.config.cluster_id:
+        skip("No cluster_id found in the environment")
+    ws.clusters.ensure_cluster_is_running(ws.config.cluster_id)
+    try:
+        # pylint: disable-next=import-outside-toplevel
+        from databricks.connect import (  # type: ignore[import-untyped]
+            DatabricksSession,
+        )
+
+        return DatabricksSession.builder.sdkConfig(ws.config).getOrCreate()
+    except ImportError:
+        skip("Please run `pip install databricks-connect`")
+        return None
diff --git a/src/databricks/labs/pytester/fixtures/plugin.py b/src/databricks/labs/pytester/fixtures/plugin.py
@@ -51,12 +51,14 @@
 from databricks.labs.pytester.fixtures.ml import make_experiment, make_model, make_serving_endpoint
 from databricks.labs.pytester.fixtures.redash import make_query
 from databricks.labs.pytester.fixtures.watchdog import watchdog_remove_after, watchdog_purge_suffix
+from databricks.labs.pytester.fixtures.connect import spark
 
 __all__ = [
     'debug_env_name',
     'debug_env',
     'env_or_skip',
     'ws',
+    'spark',
     'sql_backend',
     'sql_exec',
     'sql_fetch_all',

diff --git a/tests/integration/fixtures/test_connect.py b/tests/integration/fixtures/test_connect.py
@@ -0,0 +1,3 @@
+def test_databricks_connect(spark):
+    rows = spark.sql("SELECT 1").collect()
+    assert rows[0][0] == 1
diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py
@@ -1 +1 @@
-pytest_plugins = ['pytester']
+pytest_plugins = ['pytester']
diff --git a/tests/unit/fixtures/test_plugin.py b/tests/unit/fixtures/test_plugin.py
@@ -45,5 +45,3 @@ def test_a_thing(pytester):
     pytester.makepyfile(INLINE)
     result = pytester.runpytest()
     result.assert_outcomes(passed=1)
-
-