diff --git a/README.md b/README.md index 9ee4493..9f8cd25 100644 --- 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)] diff --git a/src/databricks/labs/pytester/fixtures/connect.py b/src/databricks/labs/pytester/fixtures/connect.py new file mode 100644 index 0000000..19b7d15 --- /dev/null +++ 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 index 5d7e2a6..ef4281b 100644 --- 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 new file mode 100644 index 0000000..53fcc8f --- /dev/null +++ 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 index b34e556..a60b5be 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -1 +1 @@ -pytest_plugins = ['pytester'] \ No newline at end of file +pytest_plugins = ['pytester'] diff --git a/tests/unit/fixtures/test_plugin.py b/tests/unit/fixtures/test_plugin.py index fd8ba38..7497e3d 100644 --- 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) - -