Added acc and make_acc_group fixtures

databrickslabs · Sep 18, 2024 · 399816d · 399816d
1 parent 2bc7842
commit 399816d
Show file tree

Hide file tree

Showing 9 changed files with 172 additions and 20 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -1,14 +1,26 @@
-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
 
 
 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()