From 399816df3cadf85b0f4c16b28610d5e93a31bec5 Mon Sep 17 00:00:00 2001 From: Serge Smertin Date: Wed, 18 Sep 2024 12:43:02 +0200 Subject: [PATCH] Added `acc` and `make_acc_group` fixtures --- README.md | 62 +++++++++++++++-- .../labs/pytester/fixtures/baseline.py | 69 +++++++++++++++---- .../labs/pytester/fixtures/environment.py | 6 ++ src/databricks/labs/pytester/fixtures/iam.py | 20 ++++++ .../labs/pytester/fixtures/plugin.py | 5 +- tests/integration/conftest.py | 7 ++ tests/integration/fixtures/test_baseline.py | 3 + tests/integration/fixtures/test_iam.py | 6 ++ tests/unit/fixtures/test_iam.py | 14 +++- 9 files changed, 172 insertions(+), 20 deletions(-) create mode 100644 tests/integration/fixtures/test_baseline.py diff --git a/README.md b/README.md index ea17cd3..e1b6f1e 100644 --- a/README.md +++ b/README.md @@ -200,7 +200,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), [`is_in_debug`](#is_in_debug-fixture). +See also [`acc`](#acc-fixture), [`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)] @@ -219,7 +219,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), [`is_in_debug`](#is_in_debug-fixture). +See also [`acc`](#acc-fixture), [`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)] @@ -245,6 +245,33 @@ def test_workspace_operations(ws): 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)] + +### `acc` fixture +Create and provide a Databricks AccountClient object. + +This fixture initializes a Databricks AccountClient object, which can be used +to interact with the Databricks account API. The created instance of AccountClient +is shared across all test functions within the test session. + +Requires `DATABRICKS_ACCOUNT_ID` environment variable to be set. If `DATABRICKS_HOST` +points to a workspace host, the fixture would automatically determine the account host +from it. + +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 AccountClient. + +In your test functions, include this fixture as an argument to use the AccountClient: + +```python +def test_listing_workspaces(acc): + workspaces = acc.workspaces.list() + assert len(workspaces) >= 1 +``` + +See also [`make_acc_group`](#make_acc_group-fixture), [`debug_env`](#debug_env-fixture), [`product_info`](#product_info-fixture), [`env_or_skip`](#env_or_skip-fixture). + + [[back to top](#python-testing-for-databricks)] ### `spark` fixture @@ -309,7 +336,7 @@ random_string = make_random(k=8) assert len(random_string) == 8 ``` -See also [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_directory`](#make_directory-fixture), [`make_experiment`](#make_experiment-fixture), [`make_group`](#make_group-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_job`](#make_job-fixture), [`make_model`](#make_model-fixture), [`make_notebook`](#make_notebook-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_query`](#make_query-fixture), [`make_repo`](#make_repo-fixture), [`make_schema`](#make_schema-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_table`](#make_table-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_warehouse`](#make_warehouse-fixture). +See also [`make_acc_group`](#make_acc_group-fixture), [`make_catalog`](#make_catalog-fixture), [`make_cluster`](#make_cluster-fixture), [`make_cluster_policy`](#make_cluster_policy-fixture), [`make_directory`](#make_directory-fixture), [`make_experiment`](#make_experiment-fixture), [`make_group`](#make_group-fixture), [`make_instance_pool`](#make_instance_pool-fixture), [`make_job`](#make_job-fixture), [`make_model`](#make_model-fixture), [`make_notebook`](#make_notebook-fixture), [`make_pipeline`](#make_pipeline-fixture), [`make_query`](#make_query-fixture), [`make_repo`](#make_repo-fixture), [`make_schema`](#make_schema-fixture), [`make_secret_scope`](#make_secret_scope-fixture), [`make_serving_endpoint`](#make_serving_endpoint-fixture), [`make_table`](#make_table-fixture), [`make_udf`](#make_udf-fixture), [`make_user`](#make_user-fixture), [`make_warehouse`](#make_warehouse-fixture). [[back to top](#python-testing-for-databricks)] @@ -504,6 +531,26 @@ def test_new_group(make_group, make_user, ws): See also [`ws`](#ws-fixture), [`make_random`](#make_random-fixture), [`watchdog_purge_suffix`](#watchdog_purge_suffix-fixture). +[[back to top](#python-testing-for-databricks)] + +### `make_acc_group` fixture +This fixture provides a function to manage Databricks account groups. Groups can be created with +specified members and roles, and they will be deleted after the test is complete. + +Has the same arguments and behavior as [`make_group` fixture](#make_group-fixture) but uses the account +client instead of the workspace client. + +Example usage: +```python +def test_new_account_group(make_acc_group, acc): + group = make_acc_group() + loaded = acc.groups.get(group.id) + assert group.display_name == loaded.display_name +``` + +See also [`acc`](#acc-fixture), [`make_random`](#make_random-fixture), [`watchdog_purge_suffix`](#watchdog_purge_suffix-fixture). + + [[back to top](#python-testing-for-databricks)] ### `make_user` fixture @@ -819,7 +866,7 @@ See also [`ws`](#ws-fixture). ### `product_info` fixture _No description yet._ -See also [`ws`](#ws-fixture). +See also [`acc`](#acc-fixture), [`ws`](#ws-fixture). [[back to top](#python-testing-for-databricks)] @@ -1020,13 +1067,16 @@ See also [`make_cluster`](#make_cluster-fixture), [`make_instance_pool`](#make_i ### `watchdog_purge_suffix` fixture 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). +See also [`make_acc_group`](#make_acc_group-fixture), [`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._ +Returns true if the test is running from a debugger in IDE, otherwise false. + +The following IDE are supported: IntelliJ IDEA (including Community Edition), +PyCharm (including Community Edition), and Visual Studio Code. See also [`debug_env`](#debug_env-fixture), [`env_or_skip`](#env_or_skip-fixture). diff --git a/src/databricks/labs/pytester/fixtures/baseline.py b/src/databricks/labs/pytester/fixtures/baseline.py index 5760514..86bf204 100644 --- a/src/databricks/labs/pytester/fixtures/baseline.py +++ b/src/databricks/labs/pytester/fixtures/baseline.py @@ -4,7 +4,8 @@ from pytest import fixture -from databricks.sdk import WorkspaceClient +from databricks.sdk import WorkspaceClient, AccountClient +from databricks.sdk.config import Config from databricks.sdk.errors import DatabricksError _LOG = logging.getLogger(__name__) @@ -138,26 +139,70 @@ def test_workspace_operations(ws): ``` """ product_name, product_version = product_info - # ignores fixed in https://github.com/databricks/databricks-sdk-py/pull/760 return WorkspaceClient( host=debug_env["DATABRICKS_HOST"], - auth_type=debug_env.get("DATABRICKS_AUTH_TYPE"), # type: ignore - token=debug_env.get("DATABRICKS_TOKEN"), # type: ignore - username=debug_env.get("DATABRICKS_USERNAME"), # type: ignore - password=debug_env.get("DATABRICKS_PASSWORD"), # type: ignore - client_id=debug_env.get("DATABRICKS_CLIENT_ID"), # type: ignore - client_secret=debug_env.get("DATABRICKS_CLIENT_SECRET"), # type: ignore + auth_type=debug_env.get("DATABRICKS_AUTH_TYPE"), + token=debug_env.get("DATABRICKS_TOKEN"), + username=debug_env.get("DATABRICKS_USERNAME"), + password=debug_env.get("DATABRICKS_PASSWORD"), + client_id=debug_env.get("DATABRICKS_CLIENT_ID"), + client_secret=debug_env.get("DATABRICKS_CLIENT_SECRET"), debug_truncate_bytes=debug_env.get("DATABRICKS_DEBUG_TRUNCATE_BYTES"), # type: ignore debug_headers=debug_env.get("DATABRICKS_DEBUG_HEADERS"), # type: ignore - azure_client_id=debug_env.get("ARM_CLIENT_ID"), # type: ignore - azure_tenant_id=debug_env.get("ARM_TENANT_ID"), # type: ignore - azure_client_secret=debug_env.get("ARM_CLIENT_SECRET"), # type: ignore - cluster_id=debug_env.get("DATABRICKS_CLUSTER_ID"), # type: ignore + azure_client_id=debug_env.get("ARM_CLIENT_ID"), + azure_tenant_id=debug_env.get("ARM_TENANT_ID"), + azure_client_secret=debug_env.get("ARM_CLIENT_SECRET"), + cluster_id=debug_env.get("DATABRICKS_CLUSTER_ID"), product=product_name, product_version=product_version, ) +@fixture +def acc(debug_env: dict[str, str], product_info: tuple[str, str], env_or_skip) -> AccountClient: + """ + Create and provide a Databricks AccountClient object. + + This fixture initializes a Databricks AccountClient object, which can be used + to interact with the Databricks account API. The created instance of AccountClient + is shared across all test functions within the test session. + + Requires `DATABRICKS_ACCOUNT_ID` environment variable to be set. If `DATABRICKS_HOST` + points to a workspace host, the fixture would automatically determine the account host + from it. + + 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 AccountClient. + + In your test functions, include this fixture as an argument to use the AccountClient: + + ```python + def test_listing_workspaces(acc): + workspaces = acc.workspaces.list() + assert len(workspaces) >= 1 + ``` + """ + product_name, product_version = product_info + config = Config( + host=debug_env["DATABRICKS_HOST"], + account_id=env_or_skip("DATABRICKS_ACCOUNT_ID"), + auth_type=debug_env.get("DATABRICKS_AUTH_TYPE"), + username=debug_env.get("DATABRICKS_USERNAME"), + password=debug_env.get("DATABRICKS_PASSWORD"), + client_id=debug_env.get("DATABRICKS_CLIENT_ID"), + client_secret=debug_env.get("DATABRICKS_CLIENT_SECRET"), + debug_truncate_bytes=debug_env.get("DATABRICKS_DEBUG_TRUNCATE_BYTES"), # type: ignore + debug_headers=debug_env.get("DATABRICKS_DEBUG_HEADERS"), # type: ignore + azure_client_id=debug_env.get("ARM_CLIENT_ID"), + azure_tenant_id=debug_env.get("ARM_TENANT_ID"), + azure_client_secret=debug_env.get("ARM_CLIENT_SECRET"), + product=product_name, + product_version=product_version, + ) + config.host = config.environment.deployment_url('accounts') + return AccountClient(config=config) + + @fixture def log_workspace_link(ws): """Returns a function to log a workspace link.""" diff --git a/src/databricks/labs/pytester/fixtures/environment.py b/src/databricks/labs/pytester/fixtures/environment.py index 2e78554..e5187e2 100644 --- a/src/databricks/labs/pytester/fixtures/environment.py +++ b/src/databricks/labs/pytester/fixtures/environment.py @@ -11,6 +11,12 @@ @fixture def is_in_debug() -> bool: + """ + Returns true if the test is running from a debugger in IDE, otherwise false. + + The following IDE are supported: IntelliJ IDEA (including Community Edition), + PyCharm (including Community Edition), and Visual Studio Code. + """ return os.path.basename(sys.argv[0]) in {"_jb_pytest_runner.py", "testlauncher.py"} diff --git a/src/databricks/labs/pytester/fixtures/iam.py b/src/databricks/labs/pytester/fixtures/iam.py index c36f055..3eb752c 100644 --- a/src/databricks/labs/pytester/fixtures/iam.py +++ b/src/databricks/labs/pytester/fixtures/iam.py @@ -70,6 +70,26 @@ def test_new_group(make_group, make_user, ws): yield from _make_group("workspace group", ws.config, ws.groups, make_random, watchdog_purge_suffix) +@fixture +def make_acc_group(acc, make_random, watchdog_purge_suffix): + """ + This fixture provides a function to manage Databricks account groups. Groups can be created with + specified members and roles, and they will be deleted after the test is complete. + + Has the same arguments and behavior as [`make_group` fixture](#make_group-fixture) but uses the account + client instead of the workspace client. + + Example usage: + ```python + def test_new_account_group(make_acc_group, acc): + group = make_acc_group() + loaded = acc.groups.get(group.id) + assert group.display_name == loaded.display_name + ``` + """ + yield from _make_group("account group", acc.config, acc.groups, make_random, watchdog_purge_suffix) + + def _scim_values(ids: list[str]) -> list[iam.ComplexValue]: return [iam.ComplexValue(value=x) for x in ids] diff --git a/src/databricks/labs/pytester/fixtures/plugin.py b/src/databricks/labs/pytester/fixtures/plugin.py index ef4281b..50c9b84 100644 --- a/src/databricks/labs/pytester/fixtures/plugin.py +++ b/src/databricks/labs/pytester/fixtures/plugin.py @@ -2,6 +2,7 @@ from databricks.labs.pytester.fixtures.baseline import ( ws, + acc, make_random, product_info, log_workspace_link, @@ -15,7 +16,7 @@ make_pipeline, make_warehouse, ) -from databricks.labs.pytester.fixtures.iam import make_group, make_user +from databricks.labs.pytester.fixtures.iam import make_group, make_acc_group, make_user from databricks.labs.pytester.fixtures.catalog import ( make_udf, make_catalog, @@ -58,6 +59,7 @@ 'debug_env', 'env_or_skip', 'ws', + 'acc', 'spark', 'sql_backend', 'sql_exec', @@ -74,6 +76,7 @@ 'make_pipeline', 'make_warehouse', 'make_group', + 'make_acc_group', 'make_user', 'make_pipeline_permissions', 'make_notebook', diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 581c857..bcf01e9 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -3,6 +3,8 @@ from pytest import fixture from databricks.labs.blueprint.logger import install_logger +from databricks.labs.pytester.__about__ import __version__ + install_logger() logging.getLogger('databricks.labs.pytester').setLevel(logging.DEBUG) @@ -11,3 +13,8 @@ @fixture def debug_env_name(): return "ucws" + + +@fixture +def product_info(): + return 'pytester', __version__ diff --git a/tests/integration/fixtures/test_baseline.py b/tests/integration/fixtures/test_baseline.py new file mode 100644 index 0000000..0c5cade --- /dev/null +++ b/tests/integration/fixtures/test_baseline.py @@ -0,0 +1,3 @@ +def test_listing_workspaces(acc): + workspaces = acc.workspaces.list() + assert len(workspaces) >= 1 diff --git a/tests/integration/fixtures/test_iam.py b/tests/integration/fixtures/test_iam.py index ece5bc0..445675c 100644 --- a/tests/integration/fixtures/test_iam.py +++ b/tests/integration/fixtures/test_iam.py @@ -13,3 +13,9 @@ def test_new_group(make_group, make_user, ws): loaded = ws.groups.get(group.id) assert group.display_name == loaded.display_name assert group.members == loaded.members + + +def test_new_account_group(make_acc_group, acc): + group = make_acc_group() + loaded = acc.groups.get(group.id) + assert group.display_name == loaded.display_name diff --git a/tests/unit/fixtures/test_iam.py b/tests/unit/fixtures/test_iam.py index ad06cd4..f8e46df 100644 --- a/tests/unit/fixtures/test_iam.py +++ b/tests/unit/fixtures/test_iam.py @@ -1,4 +1,4 @@ -from databricks.labs.pytester.fixtures.iam import make_user, make_group +from databricks.labs.pytester.fixtures.iam import make_user, make_group, make_acc_group from databricks.labs.pytester.fixtures.unwrap import call_stateful @@ -6,9 +6,21 @@ def test_make_user_no_args(): ctx, user = call_stateful(make_user) assert ctx is not None assert user is not None + ctx['ws'].users.create.assert_called_once() + ctx['ws'].users.delete.assert_called_once() def test_make_group_no_args(): ctx, group = call_stateful(make_group) assert ctx is not None assert group is not None + ctx['ws'].groups.create.assert_called_once() + ctx['ws'].groups.delete.assert_called_once() + + +def test_make_acc_group_no_args(): + ctx, group = call_stateful(make_acc_group) + assert ctx is not None + assert group is not None + ctx['acc'].groups.create.assert_called_once() + ctx['acc'].groups.delete.assert_called_once()