code review feedback

Author: mwojtyczka

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

@@ -50,7 +50,7 @@ def test_databricks_connect(spark):
         return None
 
 
-def ensure_cluster_is_running(cluster_id, ws):
+def ensure_cluster_is_running(cluster_id: str, ws: WorkspaceClient) -> None:
     if not cluster_id:
         skip("No cluster_id found in the environment")
     ws.clusters.ensure_cluster_is_running(cluster_id)

tests/integration/fixtures/test_connect.py

@@ -17,6 +17,17 @@ def debug_env_bugfix(monkeypatch, debug_env):
     monkeypatch.setitem(debug_env, "DATABRICKS_CLUSTER_ID", "1114-152544-29g1w07e")
 
 
+@fixture
+def spark_serverless_cluster_id(ws):
+    # get new spark session with serverless cluster outside the actual spark fixture under test
+    spark_serverless = DatabricksSession.builder.serverless(True).getOrCreate()
+    # get cluster id from the existing serverless spark session
+    cluster_id = spark_serverless.conf.get("spark.databricks.clusterUsageTags.clusterId")
+    ws.config.serverless_compute_id = cluster_id
+    yield cluster_id
+    spark_serverless.stop()
+
+
 def test_databricks_connect(debug_env_bugfix, ws, spark):
     rows = spark.sql("SELECT 1").collect()
     assert rows[0][0] == 1
@@ -33,20 +44,14 @@ def test_databricks_connect_serverless(serverless_env, ws, spark):
     assert not creator  # serverless clusters don't have assigned creator
 
 
-def test_databricks_connect_serverless_set_cluster_id(serverless_env, ws, request):
-    # get spark session to retrieve serverless cluster id
-    spark_serverless = DatabricksSession.builder.serverless(True).getOrCreate()
-    cluster_id = spark_serverless.conf.get("spark.databricks.clusterUsageTags.clusterId")
-    ws.config.serverless_compute_id = cluster_id
-
-    # get a new spark session for serverless from the spark fixture using provided cluster id
-    spark_serverless_new = request.getfixturevalue("spark")
-
-    rows = spark_serverless_new.sql("SELECT 1").collect()
+def test_databricks_connect_serverless_set_cluster_id(ws, spark_serverless_cluster_id, spark):
+    rows = spark.sql("SELECT 1").collect()
     assert rows[0][0] == 1
 
-    assert spark_serverless_new.conf.get("spark.databricks.clusterUsageTags.clusterId") == cluster_id
-    creator = get_cluster_creator(spark_serverless_new, ws)
+    cluster_id = spark.conf.get("spark.databricks.clusterUsageTags.clusterId")
+    assert spark_serverless_cluster_id == cluster_id
+
+    creator = get_cluster_creator(spark, ws)
     assert not creator  # serverless clusters don't have assigned creator