diff --git a/README.md b/README.md index 06788ff..17586e0 100644 --- a/README.md +++ b/README.md @@ -34,55 +34,110 @@ ## PyTest Fixtures -### `ws` fixture -Create and provide a Databricks WorkspaceClient object. - -This fixture initializes a Databricks WorkspaceClient object, which can be used -to interact with the Databricks workspace API. The created instance of WorkspaceClient -is shared across all test functions within the test session. - -See https://databricks-sdk-py.readthedocs.io/en/latest/authentication.html +### `debug_env_name` fixture +Specify the name of the debug environment. By default, it is set to `.env`, +which will try to find a [file named `.env`](https://www.dotenv.org/docs/security/env) +in any of the parent directories of the current working directory and load +the environment variables from it via the [`debug_env` fixture](#debug_env-fixture). + +Alternatively, if you are concerned of the +[risk of `.env` files getting checked into version control](https://thehackernews.com/2024/08/attackers-exploit-public-env-files-to.html), +we recommend using the `~/.databricks/debug-env.json` file to store different sets of environment variables. +The file cannot be checked into version control by design, because it is stored in the user's home directory. + +This file is used for local debugging and integration tests in IDEs like PyCharm, VSCode, and IntelliJ IDEA +while developing Databricks Platform Automation Stack, which includes Databricks SDKs for Python, Go, and Java, +as well as Databricks Terraform Provider and Databricks CLI. This file enables multi-environment and multi-cloud +testing with a single set of integration tests. + +The file is typically structured as follows: + +```shell +$ cat ~/.databricks/debug-env.json +{ + "ws": { + "CLOUD_ENV": "azure", + "DATABRICKS_HOST": "....azuredatabricks.net", + "DATABRICKS_CLUSTER_ID": "0708-200540-...", + "DATABRICKS_WAREHOUSE_ID": "33aef...", + ... + }, + "acc": { + "CLOUD_ENV": "aws", + "DATABRICKS_HOST": "accounts.cloud.databricks.net", + "DATABRICKS_CLIENT_ID": "....", + "DATABRICKS_CLIENT_SECRET": "....", + ... + } +} +``` -Returns: --------- -databricks.sdk.WorkspaceClient: - An instance of WorkspaceClient for interacting with the Databricks Workspace APIs. +And you can load it in your `conftest.py` file as follows: -Usage: ------- -In your test functions, include this fixture as an argument to use the WorkspaceClient: +```python +@pytest.fixture +def debug_env_name(): + return "ws" +``` -.. code-block:: python +This will load the `ws` environment from the `~/.databricks/debug-env.json` file. - def test_workspace_operations(ws): - clusters = ws.clusters.list_clusters() - assert len(clusters) >= 0 +If any of the environment variables are not found, [`env_or_skip` fixture](#env_or_skip-fixture) +will gracefully skip the execution of tests. -This fixture is built on top of: [`debug_env`](#debug_env-fixture), [`product_info`](#product_info-fixture) +See also [`debug_env` fixture](#debug_env-fixture). [[back to top](#python-testing-for-databricks)] ### `debug_env` fixture -_No description yet._ +Loads environment variables specified in [`debug_env_name` fixture](#debug_env_name-fixture) from a file +for local debugging in IDEs, otherwise allowing the tests to run with the default environment variables +specified in the CI/CD pipeline. -This fixture is built on top of: [`monkeypatch`](#monkeypatch-fixture), [`debug_env_name`](#debug_env_name-fixture) +See also [`env_or_skip` fixture](#env_or_skip-fixture), [`ws` fixture](#ws-fixture), [`debug_env_name` fixture](#debug_env_name-fixture). [[back to top](#python-testing-for-databricks)] -### `debug_env_name` fixture -_No description yet._ +### `env_or_skip` fixture +Fixture to get environment variables or skip tests. -This fixture is built on top of: +It is extremely useful to skip tests if the required environment variables are not set. + +In the following example, `test_something` would only run if the environment variable +`SOME_EXTERNAL_SERVICE_TOKEN` is set: + +```python +def test_something(env_or_skip): + token = env_or_skip("SOME_EXTERNAL_SERVICE_TOKEN") + assert token is not None +``` + +See also [`make_udf` fixture](#make_udf-fixture), [`sql_backend` fixture](#sql_backend-fixture), [`debug_env` fixture](#debug_env-fixture). [[back to top](#python-testing-for-databricks)] -### `env_or_skip` fixture -_No description yet._ +### `ws` fixture +Create and provide a Databricks WorkspaceClient object. -This fixture is built on top of: [`debug_env`](#debug_env-fixture) +This fixture initializes a Databricks WorkspaceClient object, which can be used +to interact with the Databricks workspace API. The created instance of WorkspaceClient +is shared across all test functions within the test session. + +See [detailed documentation](https://databricks-sdk-py.readthedocs.io/en/latest/authentication.html) for the list +of environment variables that can be used to authenticate the WorkspaceClient. + +In your test functions, include this fixture as an argument to use the WorkspaceClient: + +```python +def test_workspace_operations(ws): + clusters = ws.clusters.list_clusters() + assert len(clusters) >= 0 +``` + +See also [`make_catalog` fixture](#make_catalog-fixture), [`make_cluster` fixture](#make_cluster-fixture), [`make_cluster_policy` fixture](#make_cluster_policy-fixture), [`make_directory` fixture](#make_directory-fixture), [`make_group` fixture](#make_group-fixture), [`make_instance_pool` fixture](#make_instance_pool-fixture), [`make_job` fixture](#make_job-fixture), [`make_notebook` fixture](#make_notebook-fixture), [`make_repo` fixture](#make_repo-fixture), [`make_schema` fixture](#make_schema-fixture), [`make_secret_scope` fixture](#make_secret_scope-fixture), [`make_secret_scope_acl` fixture](#make_secret_scope_acl-fixture), [`make_table` fixture](#make_table-fixture), [`make_udf` fixture](#make_udf-fixture), [`make_user` fixture](#make_user-fixture), [`sql_backend` fixture](#sql_backend-fixture), [`workspace_library` fixture](#workspace_library-fixture), [`debug_env` fixture](#debug_env-fixture), [`product_info` fixture](#product_info-fixture). [[back to top](#python-testing-for-databricks)] @@ -108,7 +163,7 @@ random_string = make_random(k=8) assert len(random_string) == 8 ``` -This fixture is built on top of: +See also [`make_catalog` fixture](#make_catalog-fixture), [`make_cluster` fixture](#make_cluster-fixture), [`make_cluster_policy` fixture](#make_cluster_policy-fixture), [`make_directory` fixture](#make_directory-fixture), [`make_group` fixture](#make_group-fixture), [`make_instance_pool` fixture](#make_instance_pool-fixture), [`make_job` fixture](#make_job-fixture), [`make_notebook` fixture](#make_notebook-fixture), [`make_repo` fixture](#make_repo-fixture), [`make_schema` fixture](#make_schema-fixture), [`make_secret_scope` fixture](#make_secret_scope-fixture), [`make_table` fixture](#make_table-fixture), [`make_udf` fixture](#make_udf-fixture), [`make_user` fixture](#make_user-fixture), [`workspace_library` fixture](#workspace_library-fixture). [[back to top](#python-testing-for-databricks)] @@ -141,7 +196,7 @@ To manage Databricks instance pools using the make_instance_pool fixture: instance_pool_info = make_instance_pool(instance_pool_name="my-pool") assert instance_pool_info is not None -This fixture is built on top of: [`ws`](#ws-fixture), [`make_random`](#make_random-fixture) +See also [`ws` fixture](#ws-fixture), [`make_random` fixture](#make_random-fixture). [[back to top](#python-testing-for-databricks)] @@ -176,7 +231,7 @@ To manage Databricks jobs using the make_job fixture: job_info = make_job(name="my-job") assert job_info is not None -This fixture is built on top of: [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`make_notebook`](#make_notebook-fixture) +See also [`ws` fixture](#ws-fixture), [`make_random` fixture](#make_random-fixture), [`make_notebook` fixture](#make_notebook-fixture). [[back to top](#python-testing-for-databricks)] @@ -209,7 +264,7 @@ To manage Databricks clusters using the make_cluster fixture: cluster_info = make_cluster(cluster_name="my-cluster", single_node=True) assert cluster_info is not None -This fixture is built on top of: [`ws`](#ws-fixture), [`make_random`](#make_random-fixture) +See also [`ws` fixture](#ws-fixture), [`make_random` fixture](#make_random-fixture). [[back to top](#python-testing-for-databricks)] @@ -242,7 +297,7 @@ To manage Databricks cluster policies using the make_cluster_policy fixture: policy_info = make_cluster_policy(name="my-policy") assert policy_info is not None -This fixture is built on top of: [`ws`](#ws-fixture), [`make_random`](#make_random-fixture) +See also [`ws` fixture](#ws-fixture), [`make_random` fixture](#make_random-fixture). [[back to top](#python-testing-for-databricks)] @@ -275,7 +330,7 @@ To manage Databricks workspace groups using the make_group fixture: group_info = make_group(members=["user@example.com"], roles=["viewer"]) assert group_info is not None -This fixture is built on top of: [`ws`](#ws-fixture), [`make_random`](#make_random-fixture) +See also [`ws` fixture](#ws-fixture), [`make_random` fixture](#make_random-fixture). [[back to top](#python-testing-for-databricks)] @@ -308,7 +363,7 @@ To manage Databricks workspace users using the make_user fixture: user_info = make_user() assert user_info is not None -This fixture is built on top of: [`ws`](#ws-fixture), [`make_random`](#make_random-fixture) +See also [`ws` fixture](#ws-fixture), [`make_random` fixture](#make_random-fixture). [[back to top](#python-testing-for-databricks)] @@ -341,7 +396,7 @@ To manage Databricks notebooks using the make_notebook fixture: notebook_path = make_notebook() assert notebook_path.startswith("/Users/") and notebook_path.endswith(".py") -This fixture is built on top of: [`ws`](#ws-fixture), [`make_random`](#make_random-fixture) +See also [`make_job` fixture](#make_job-fixture), [`ws` fixture](#ws-fixture), [`make_random` fixture](#make_random-fixture). [[back to top](#python-testing-for-databricks)] @@ -374,7 +429,7 @@ To manage Databricks directories using the make_directory fixture: directory_path = make_directory() assert directory_path.startswith("/Users/") and not directory_path.endswith(".py") -This fixture is built on top of: [`ws`](#ws-fixture), [`make_random`](#make_random-fixture) +See also [`ws` fixture](#ws-fixture), [`make_random` fixture](#make_random-fixture). [[back to top](#python-testing-for-databricks)] @@ -407,7 +462,7 @@ To manage Databricks repos using the make_repo fixture: repo_info = make_repo() assert repo_info is not None -This fixture is built on top of: [`ws`](#ws-fixture), [`make_random`](#make_random-fixture) +See also [`ws` fixture](#ws-fixture), [`make_random` fixture](#make_random-fixture). [[back to top](#python-testing-for-databricks)] @@ -441,7 +496,7 @@ To create a secret scope and use it within a test function: secret_scope_name = make_secret_scope() assert secret_scope_name.startswith("sdk-") -This fixture is built on top of: [`ws`](#ws-fixture), [`make_random`](#make_random-fixture) +See also [`ws` fixture](#ws-fixture), [`make_random` fixture](#make_random-fixture). [[back to top](#python-testing-for-databricks)] @@ -477,7 +532,7 @@ To manage secret scope ACLs using the make_secret_scope_acl fixture: acl_info = make_secret_scope_acl(scope=scope_name, principal=principal_name, permission=permission) assert acl_info == (scope_name, principal_name) -This fixture is built on top of: [`ws`](#ws-fixture) +See also [`ws` fixture](#ws-fixture). [[back to top](#python-testing-for-databricks)] @@ -485,7 +540,7 @@ This fixture is built on top of: [`ws`](#ws-fixture) ### `make_udf` fixture _No description yet._ -This fixture is built on top of: [`ws`](#ws-fixture), [`env_or_skip`](#env_or_skip-fixture), [`sql_backend`](#sql_backend-fixture), [`make_schema`](#make_schema-fixture), [`make_random`](#make_random-fixture) +See also [`ws` fixture](#ws-fixture), [`env_or_skip` fixture](#env_or_skip-fixture), [`sql_backend` fixture](#sql_backend-fixture), [`make_schema` fixture](#make_schema-fixture), [`make_random` fixture](#make_random-fixture). [[back to top](#python-testing-for-databricks)] @@ -493,7 +548,7 @@ This fixture is built on top of: [`ws`](#ws-fixture), [`env_or_skip`](#env_or_sk ### `make_catalog` fixture _No description yet._ -This fixture is built on top of: [`ws`](#ws-fixture), [`sql_backend`](#sql_backend-fixture), [`make_random`](#make_random-fixture) +See also [`ws` fixture](#ws-fixture), [`sql_backend` fixture](#sql_backend-fixture), [`make_random` fixture](#make_random-fixture). [[back to top](#python-testing-for-databricks)] @@ -501,7 +556,7 @@ This fixture is built on top of: [`ws`](#ws-fixture), [`sql_backend`](#sql_backe ### `make_schema` fixture _No description yet._ -This fixture is built on top of: [`ws`](#ws-fixture), [`sql_backend`](#sql_backend-fixture), [`make_random`](#make_random-fixture) +See also [`make_table` fixture](#make_table-fixture), [`make_udf` fixture](#make_udf-fixture), [`ws` fixture](#ws-fixture), [`sql_backend` fixture](#sql_backend-fixture), [`make_random` fixture](#make_random-fixture). [[back to top](#python-testing-for-databricks)] @@ -509,7 +564,7 @@ This fixture is built on top of: [`ws`](#ws-fixture), [`sql_backend`](#sql_backe ### `make_table` fixture _No description yet._ -This fixture is built on top of: [`ws`](#ws-fixture), [`sql_backend`](#sql_backend-fixture), [`make_schema`](#make_schema-fixture), [`make_random`](#make_random-fixture) +See also [`ws` fixture](#ws-fixture), [`sql_backend` fixture](#sql_backend-fixture), [`make_schema` fixture](#make_schema-fixture), [`make_random` fixture](#make_random-fixture). [[back to top](#python-testing-for-databricks)] @@ -517,7 +572,7 @@ This fixture is built on top of: [`ws`](#ws-fixture), [`sql_backend`](#sql_backe ### `product_info` fixture _No description yet._ -This fixture is built on top of: +See also [`ws` fixture](#ws-fixture). [[back to top](#python-testing-for-databricks)] @@ -525,7 +580,9 @@ This fixture is built on top of: ### `sql_backend` fixture te and provide a SQL backend for executing statements. -This fixture is built on top of: [`ws`](#ws-fixture), [`env_or_skip`](#env_or_skip-fixture) +Requires the environment variable `DATABRICKS_WAREHOUSE_ID` to be set. + +See also [`make_catalog` fixture](#make_catalog-fixture), [`make_schema` fixture](#make_schema-fixture), [`make_table` fixture](#make_table-fixture), [`make_udf` fixture](#make_udf-fixture), [`sql_exec` fixture](#sql_exec-fixture), [`sql_fetch_all` fixture](#sql_fetch_all-fixture), [`ws` fixture](#ws-fixture), [`env_or_skip` fixture](#env_or_skip-fixture). [[back to top](#python-testing-for-databricks)] @@ -533,7 +590,7 @@ This fixture is built on top of: [`ws`](#ws-fixture), [`env_or_skip`](#env_or_sk ### `sql_exec` fixture ute SQL statement and don't return any results. -This fixture is built on top of: [`sql_backend`](#sql_backend-fixture) +See also [`sql_backend` fixture](#sql_backend-fixture). [[back to top](#python-testing-for-databricks)] @@ -541,7 +598,7 @@ This fixture is built on top of: [`sql_backend`](#sql_backend-fixture) ### `sql_fetch_all` fixture h all rows from a SQL statement. -This fixture is built on top of: [`sql_backend`](#sql_backend-fixture) +See also [`sql_backend` fixture](#sql_backend-fixture). [[back to top](#python-testing-for-databricks)] @@ -549,7 +606,7 @@ This fixture is built on top of: [`sql_backend`](#sql_backend-fixture) ### `workspace_library` fixture _No description yet._ -This fixture is built on top of: [`ws`](#ws-fixture), [`fresh_local_wheel_file`](#fresh_local_wheel_file-fixture), [`make_random`](#make_random-fixture) +See also [`ws` fixture](#ws-fixture), [`make_random` fixture](#make_random-fixture). [[back to top](#python-testing-for-databricks)] diff --git a/scripts/gen-readme.py b/scripts/gen-readme.py index fb87057..0e46459 100644 --- a/scripts/gen-readme.py +++ b/scripts/gen-readme.py @@ -1,3 +1,4 @@ +import collections import inspect import re from dataclasses import dataclass @@ -19,11 +20,11 @@ def main(): class Fixture: name: str description: str - upstreams: list[str] + see_also: list[str] @staticmethod def ref(name: str) -> str: - return f"[`{name}`](#{name}-fixture)" + return f"[`{name}` fixture](#{name}-fixture)" def usage(self) -> str: lines = "\n".join(_[4:] for _ in self.description.split("\n")) @@ -35,7 +36,7 @@ def doc(self) -> str: f"### `{self.name}` fixture", self.usage() if self.description else "_No description yet._", "", - f"This fixture is built on top of: {', '.join([self.ref(up) for up in self.upstreams])}", + f"See also {', '.join([self.ref(up) for up in self.see_also])}.", "", BACK_TO_TOP, ] @@ -54,18 +55,27 @@ def overwrite_readme(part, docs): def discover_fixtures(): fixtures = [] + see_also = collections.defaultdict(set) + idx: dict[str, int] = {} for fixture in P.__all__: fn = getattr(P, fixture) upstreams = [] sig = inspect.signature(fn) for param in sig.parameters.values(): + if param.name in {'fresh_local_wheel_file', 'monkeypatch'}: + continue upstreams.append(param.name) + see_also[param.name].add(fixture) fx = Fixture( name=fixture, description=fn.__doc__, - upstreams=upstreams, + see_also=upstreams, ) + idx[fixture] = len(fixtures) fixtures.append(fx) + for fixture, other in see_also.items(): + fx = fixtures[idx[fixture]] + fx.see_also = sorted(other) + fx.see_also return fixtures diff --git a/src/databricks/labs/pytester/fixtures/baseline.py b/src/databricks/labs/pytester/fixtures/baseline.py index 0299b3e..92ee3e0 100644 --- a/src/databricks/labs/pytester/fixtures/baseline.py +++ b/src/databricks/labs/pytester/fixtures/baseline.py @@ -1,21 +1,16 @@ -import json import logging -import os -import pathlib import random import string -import sys -from collections.abc import Callable, MutableMapping from datetime import timedelta, datetime, timezone from functools import partial -import pytest from pytest import fixture from databricks.labs.lsql.backends import StatementExecutionBackend from databricks.sdk import WorkspaceClient from databricks.sdk.errors import NotFound + _LOG = logging.getLogger(__name__) @@ -67,53 +62,6 @@ def inner(k=16) -> str: return inner -def _is_in_debug() -> bool: - return os.path.basename(sys.argv[0]) in {"_jb_pytest_runner.py", "testlauncher.py"} - - -@fixture -def debug_env_name(): - # Alternatively, we could use @pytest.mark.xxx, but - # not sure how reusable it becomes then. - # - # we don't use scope=session, as monkeypatch.setenv - # doesn't work on a session level - return "UNKNOWN" - - -@pytest.fixture -def debug_env(monkeypatch, debug_env_name) -> MutableMapping[str, str]: - if not _is_in_debug(): - return os.environ - # TODO: add support for `.env` files - conf_file = pathlib.Path.home() / ".databricks/debug-env.json" - if not conf_file.exists(): - return os.environ - with conf_file.open("r") as f: - conf = json.load(f) - if debug_env_name not in conf: - sys.stderr.write(f"""{debug_env_name} not found in ~/.databricks/debug-env.json""") - msg = f"{debug_env_name} not found in ~/.databricks/debug-env.json" - raise KeyError(msg) - for env_key, value in conf[debug_env_name].items(): - monkeypatch.setenv(env_key, value) - return os.environ - - -@fixture -def env_or_skip(debug_env) -> Callable[[str], str]: - skip = pytest.skip - if _is_in_debug(): - skip = pytest.fail # type: ignore[assignment] - - def inner(var: str) -> str: - if var not in debug_env: - skip(f"Environment variable {var} is missing") - return debug_env[var] - - return inner - - def factory(name, create, remove): """ Factory function for creating fixtures. @@ -186,22 +134,16 @@ def ws(debug_env, product_info) -> WorkspaceClient: to interact with the Databricks workspace API. The created instance of WorkspaceClient is shared across all test functions within the test session. - See https://databricks-sdk-py.readthedocs.io/en/latest/authentication.html + See [detailed documentation](https://databricks-sdk-py.readthedocs.io/en/latest/authentication.html) for the list + of environment variables that can be used to authenticate the WorkspaceClient. - Returns: - -------- - databricks.sdk.WorkspaceClient: - An instance of WorkspaceClient for interacting with the Databricks Workspace APIs. - - Usage: - ------ In your test functions, include this fixture as an argument to use the WorkspaceClient: - .. code-block:: python - - def test_workspace_operations(ws): - clusters = ws.clusters.list_clusters() - assert len(clusters) >= 0 + ```python + def test_workspace_operations(ws): + clusters = ws.clusters.list_clusters() + assert len(clusters) >= 0 + ``` """ product_name, product_version = product_info return WorkspaceClient(host=debug_env["DATABRICKS_HOST"], product=product_name, product_version=product_version) @@ -209,7 +151,10 @@ def test_workspace_operations(ws): @fixture def sql_backend(ws, env_or_skip) -> StatementExecutionBackend: - """Create and provide a SQL backend for executing statements.""" + """Create and provide a SQL backend for executing statements. + + Requires the environment variable `DATABRICKS_WAREHOUSE_ID` to be set. + """ warehouse_id = env_or_skip("DATABRICKS_WAREHOUSE_ID") return StatementExecutionBackend(ws, warehouse_id) diff --git a/src/databricks/labs/pytester/fixtures/environment.py b/src/databricks/labs/pytester/fixtures/environment.py new file mode 100644 index 0000000..5c9d375 --- /dev/null +++ b/src/databricks/labs/pytester/fixtures/environment.py @@ -0,0 +1,147 @@ +import json +import os +from pathlib import Path +import sys +from collections.abc import MutableMapping, Callable + +import pytest +from pytest import fixture +from databricks.labs.blueprint.entrypoint import find_dir_with_leaf + + +def _is_in_debug() -> bool: + return os.path.basename(sys.argv[0]) in {"_jb_pytest_runner.py", "testlauncher.py"} + + +@fixture +def debug_env_name(): + """ + Specify the name of the debug environment. By default, it is set to `.env`, + which will try to find a [file named `.env`](https://www.dotenv.org/docs/security/env) + in any of the parent directories of the current working directory and load + the environment variables from it via the [`debug_env` fixture](#debug_env-fixture). + + Alternatively, if you are concerned of the + [risk of `.env` files getting checked into version control](https://thehackernews.com/2024/08/attackers-exploit-public-env-files-to.html), + we recommend using the `~/.databricks/debug-env.json` file to store different sets of environment variables. + The file cannot be checked into version control by design, because it is stored in the user's home directory. + + This file is used for local debugging and integration tests in IDEs like PyCharm, VSCode, and IntelliJ IDEA + while developing Databricks Platform Automation Stack, which includes Databricks SDKs for Python, Go, and Java, + as well as Databricks Terraform Provider and Databricks CLI. This file enables multi-environment and multi-cloud + testing with a single set of integration tests. + + The file is typically structured as follows: + + ```shell + $ cat ~/.databricks/debug-env.json + { + "ws": { + "CLOUD_ENV": "azure", + "DATABRICKS_HOST": "....azuredatabricks.net", + "DATABRICKS_CLUSTER_ID": "0708-200540-...", + "DATABRICKS_WAREHOUSE_ID": "33aef...", + ... + }, + "acc": { + "CLOUD_ENV": "aws", + "DATABRICKS_HOST": "accounts.cloud.databricks.net", + "DATABRICKS_CLIENT_ID": "....", + "DATABRICKS_CLIENT_SECRET": "....", + ... + } + } + ``` + + And you can load it in your `conftest.py` file as follows: + + ```python + @pytest.fixture + def debug_env_name(): + return "ws" + ``` + + This will load the `ws` environment from the `~/.databricks/debug-env.json` file. + + If any of the environment variables are not found, [`env_or_skip` fixture](#env_or_skip-fixture) + will gracefully skip the execution of tests. + """ + return ".env" + + +@fixture +def debug_env(monkeypatch, debug_env_name) -> MutableMapping[str, str]: + """ + Loads environment variables specified in [`debug_env_name` fixture](#debug_env_name-fixture) from a file + for local debugging in IDEs, otherwise allowing the tests to run with the default environment variables + specified in the CI/CD pipeline. + """ + if not _is_in_debug(): + return os.environ + if debug_env_name == ".env": + dot_env = _parse_dotenv() + if not dot_env: + return os.environ + for env_key, value in dot_env.items(): + monkeypatch.setenv(env_key, value) + return os.environ + conf_file = Path.home() / ".databricks/debug-env.json" + if not conf_file.exists(): + return os.environ + with conf_file.open("r") as f: + conf = json.load(f) + if debug_env_name not in conf: + sys.stderr.write(f"""{debug_env_name} not found in ~/.databricks/debug-env.json""") + msg = f"{debug_env_name} not found in ~/.databricks/debug-env.json" + raise KeyError(msg) + for env_key, value in conf[debug_env_name].items(): + monkeypatch.setenv(env_key, value) + return os.environ + + +@fixture +def env_or_skip(debug_env) -> Callable[[str], str]: + """ + Fixture to get environment variables or skip tests. + + It is extremely useful to skip tests if the required environment variables are not set. + + In the following example, `test_something` would only run if the environment variable + `SOME_EXTERNAL_SERVICE_TOKEN` is set: + + ```python + def test_something(env_or_skip): + token = env_or_skip("SOME_EXTERNAL_SERVICE_TOKEN") + assert token is not None + ``` + """ + skip = pytest.skip + if _is_in_debug(): + skip = pytest.fail # type: ignore[assignment] + + def inner(var: str) -> str: + if var not in debug_env: + skip(f"Environment variable {var} is missing") + return debug_env[var] + + return inner + + +def _parse_dotenv(): + """See https://www.dotenv.org/docs/security/env""" + dot_env = find_dir_with_leaf(Path.cwd(), '.env') + if dot_env is None: + return {} + env_vars = {} + with dot_env.open(encoding='utf8') as file: + for line in file: + line = line.strip() + if not line or line.startswith('#'): + continue + if '=' not in line: + continue + key, value = line.split('=', 1) + # Remove any surrounding quotes (single or double) + value = value.strip().strip('"').strip("'") + env_vars[key.strip()] = value + return env_vars diff --git a/src/databricks/labs/pytester/fixtures/plugin.py b/src/databricks/labs/pytester/fixtures/plugin.py index dc5a8fe..26d9861 100644 --- a/src/databricks/labs/pytester/fixtures/plugin.py +++ b/src/databricks/labs/pytester/fixtures/plugin.py @@ -6,9 +6,6 @@ sql_backend, sql_exec, sql_fetch_all, - env_or_skip, - debug_env, - debug_env_name, product_info, ) from databricks.labs.pytester.fixtures.compute import make_instance_pool, make_job, make_cluster, make_cluster_policy @@ -17,12 +14,13 @@ from databricks.labs.pytester.fixtures.notebooks import make_notebook, make_directory, make_repo from databricks.labs.pytester.fixtures.secrets import make_secret_scope, make_secret_scope_acl from databricks.labs.pytester.fixtures.wheel import workspace_library +from databricks.labs.pytester.fixtures.environment import debug_env, debug_env_name, env_or_skip __all__ = [ - 'ws', - 'debug_env', 'debug_env_name', + 'debug_env', 'env_or_skip', + 'ws', 'make_random', 'make_instance_pool', 'make_job',