Merge branch 'main' into feat/add-file-option-to-make-job

Author: JCZuurmond

src/databricks/labs/pytester/fixtures/baseline.py

@@ -1,6 +1,8 @@
 import logging
 import random
 import string
+from collections.abc import Callable, Generator
+from typing import Any, TypeVar
 
 from pytest import fixture
 
@@ -55,7 +57,14 @@ def inner(k=16) -> str:
     return inner
 
 
-def factory(name, create, remove):
+T = TypeVar("T")
+
+
+def factory(
+    name: str,
+    create: Callable[..., T],
+    remove: Callable[[T], None],
+) -> Generator[Callable[..., T], None, None]:
     """
     Factory function for creating fixtures.
 
@@ -95,13 +104,13 @@ def cleanup(scope):
 
            yield from factory("secret scope", create, cleanup)
     """
-    cleanup = []
+    cleanup: list[T] = []
 
-    def inner(**kwargs):
-        some = create(**kwargs)
-        _LOG.debug(f"added {name} fixture: {some}")
-        cleanup.append(some)
-        return some
+    def inner(**kwargs: Any) -> T:
+        out = create(**kwargs)
+        _LOG.debug(f"added {name} fixture: {out}")
+        cleanup.append(out)
+        return out
 
     yield inner
     _LOG.debug(f"clearing {len(cleanup)} {name} fixtures")

src/databricks/labs/pytester/fixtures/compute.py

@@ -1,6 +1,6 @@
 import json
 import warnings
-from collections.abc import Generator
+from collections.abc import Callable, Generator
 from pathlib import Path
 from unittest.mock import Mock
 
@@ -31,7 +31,7 @@ def make_cluster_policy(
     make_random,
     log_workspace_link,
     watchdog_purge_suffix,
-) -> Generator[CreatePolicyResponse, None, None]:
+) -> Generator[Callable[..., CreatePolicyResponse], None, None]:
     """
     Create a Databricks cluster policy and clean it up after the test. Returns a function to create cluster policies,
     which returns `databricks.sdk.service.compute.CreatePolicyResponse` instance.
@@ -63,7 +63,9 @@ def create(*, name: str | None = None, **kwargs) -> CreatePolicyResponse:
 
 
 @fixture
-def make_cluster(ws, make_random, log_workspace_link, watchdog_remove_after) -> Generator[ClusterDetails, None, None]:
+def make_cluster(
+    ws, make_random, log_workspace_link, watchdog_remove_after
+) -> Generator[Callable[..., Wait[ClusterDetails]], None, None]:
     """
     Create a Databricks cluster, waits for it to start, and clean it up after the test.
     Returns a function to create clusters. You can get `cluster_id` attribute from the returned object.
@@ -127,7 +129,7 @@ def make_instance_pool(
     make_random,
     log_workspace_link,
     watchdog_remove_after,
-) -> Generator[CreateInstancePoolResponse, None, None]:
+) -> Generator[Callable[..., CreateInstancePoolResponse], None, None]:
     """
     Create a Databricks instance pool and clean it up after the test. Returns a function to create instance pools.
     Use `instance_pool_id` attribute from the returned object to get an ID of the pool.
@@ -169,7 +171,7 @@ def make_job(
     make_workspace_file,
     log_workspace_link,
     watchdog_remove_after,
-) -> Generator[Job, None, None]:
+) -> Generator[Callable[..., Job], None, None]:
     """
     Create a Databricks job and clean it up after the test. Returns a function to create jobs, that returns
     a `databricks.sdk.service.jobs.Job` instance.
@@ -265,7 +267,7 @@ def make_pipeline(
     make_notebook,
     watchdog_remove_after,
     watchdog_purge_suffix,
-) -> Generator[CreatePipelineResponse, None, None]:
+) -> Generator[Callable[..., CreatePipelineResponse], None, None]:
     """
     Create Delta Live Table Pipeline and clean it up after the test. Returns a function to create pipelines.
     Results in a `databricks.sdk.service.pipelines.CreatePipelineResponse` instance.
@@ -309,7 +311,9 @@ def create(**kwargs) -> CreatePipelineResponse:
 
 
 @fixture
-def make_warehouse(ws, make_random, watchdog_remove_after) -> Generator[Wait[GetWarehouseResponse], None, None]:
+def make_warehouse(
+    ws, make_random, watchdog_remove_after
+) -> Generator[Callable[..., Wait[GetWarehouseResponse]], None, None]:
     """
     Create a Databricks warehouse and clean it up after the test. Returns a function to create warehouses.
 

src/databricks/labs/pytester/fixtures/iam.py

@@ -1,6 +1,6 @@
 import logging
 import warnings
-from collections.abc import Generator
+from collections.abc import Callable, Generator
 from datetime import timedelta
 
 from pytest import fixture
@@ -140,7 +140,9 @@ def _double_check_group_in_listing() -> None:
     _double_check_group_in_listing()
 
 
-def _make_group(name: str, cfg: Config, interface, make_random, watchdog_purge_suffix) -> Generator[Group, None, None]:
+def _make_group(
+    name: str, cfg: Config, interface, make_random, watchdog_purge_suffix
+) -> Generator[Callable[..., Group], None, None]:
     _not_specified = object()
 
     @retried(on=[ResourceConflict], timeout=timedelta(seconds=30))

src/databricks/labs/pytester/fixtures/ml.py

@@ -1,4 +1,4 @@
-from collections.abc import Generator, Callable
+from collections.abc import Callable, Generator
 
 from pytest import fixture
 from databricks.sdk.service._internal import Wait
@@ -9,7 +9,7 @@
     ServedModelInputWorkloadSize,
     EndpointTag,
 )
-from databricks.sdk.service.ml import CreateExperimentResponse, ModelTag, GetModelResponse
+from databricks.sdk.service.ml import CreateExperimentResponse, ModelDatabricks, ModelTag
 
 from databricks.labs.pytester.fixtures.baseline import factory
 
@@ -21,7 +21,7 @@ def make_experiment(
     make_directory,
     log_workspace_link,
     watchdog_purge_suffix,
-) -> Generator[CreateExperimentResponse, None, None]:
+) -> Generator[Callable[..., CreateExperimentResponse], None, None]:
     """
     Returns a function to create Databricks Experiments and clean them up after the test.
     The function returns a `databricks.sdk.service.ml.CreateExperimentResponse` object.
@@ -63,7 +63,7 @@ def create(
 
 
 @fixture
-def make_model(ws, make_random, watchdog_remove_after) -> Generator[Callable[..., GetModelResponse], None, None]:
+def make_model(ws, make_random, watchdog_remove_after) -> Generator[Callable[..., ModelDatabricks], None, None]:
     """
     Returns a function to create Databricks Models and clean them up after the test.
     The function returns a `databricks.sdk.service.ml.GetModelResponse` object.
@@ -86,7 +86,7 @@ def test_models(make_group, make_model, make_registered_model_permissions):
     ```
     """
 
-    def create(*, model_name: str | None = None, **kwargs) -> GetModelResponse:
+    def create(*, model_name: str | None = None, **kwargs) -> ModelDatabricks:
         if model_name is None:
             model_name = f"dummy-{make_random(4)}"
         remove_after_tag = ModelTag(key="RemoveAfter", value=watchdog_remove_after)
@@ -96,6 +96,7 @@ def create(*, model_name: str | None = None, **kwargs) -> GetModelResponse:
             kwargs["tags"].append(remove_after_tag)
         created_model = ws.model_registry.create_model(model_name, **kwargs)
         model = ws.model_registry.get_model(created_model.registered_model.name)
+        assert model.registered_model_databricks is not None
         return model.registered_model_databricks
 
     yield from factory("model", create, lambda item: ws.model_registry.delete_model(item.id))

src/databricks/labs/pytester/fixtures/notebooks.py

@@ -1,7 +1,7 @@
 import io
 import logging
 import sys
-from collections.abc import Generator, Callable
+from collections.abc import Callable, Generator
 from pathlib import Path
 from unittest.mock import Mock
 
@@ -128,7 +128,9 @@ def create(
 
 
 @fixture
-def make_directory(ws: WorkspaceClient, make_random, watchdog_purge_suffix) -> Generator[WorkspacePath, None, None]:
+def make_directory(
+    ws: WorkspaceClient, make_random, watchdog_purge_suffix
+) -> Generator[Callable[..., WorkspacePath], None, None]:
     """
     Returns a function to create Databricks Workspace Folders and clean them up after the test.
     The function returns [`os.PathLike` object](https://github.com/databrickslabs/blueprint?tab=readme-ov-file#python-native-pathlibpath-like-interfaces).
@@ -159,7 +161,7 @@ def create(*, path: str | Path | None = None) -> WorkspacePath:
 
 
 @fixture
-def make_repo(ws, make_random, watchdog_purge_suffix) -> Generator[RepoInfo, None, None]:
+def make_repo(ws, make_random, watchdog_purge_suffix) -> Generator[Callable[..., RepoInfo], None, None]:
     """
     Returns a function to create Databricks Repos and clean them up after the test.
     The function returns a `databricks.sdk.service.workspace.RepoInfo` object.

src/databricks/labs/pytester/fixtures/redash.py

@@ -1,5 +1,5 @@
 import json
-from collections.abc import Generator
+from collections.abc import Callable, Generator
 
 from pytest import fixture
 from databricks.sdk.service.sql import LegacyQuery
@@ -14,7 +14,7 @@ def make_query(
     make_random,
     log_workspace_link,
     watchdog_remove_after,
-) -> Generator[LegacyQuery, None, None]:
+) -> Generator[Callable[..., LegacyQuery], None, None]:
     """
     Create a query and remove it after the test is done. Returns the `databricks.sdk.service.sql.LegacyQuery` object.