Integration test changed, unit tests

src/databricks/labs/ucx/install.py

@@ -262,7 +262,7 @@ def run_for_config(
     ) -> "WorkspaceInstaller":
         workspace_installer = WorkspaceInstaller(ws, prefix=prefix, wheels=wheels, sql_backend=sql_backend)
         logger.info(f"Installing UCX v{workspace_installer._product_info.version()} on {ws.config.host}")
-        workspace_installer._enable_files_in_repos()
+        workspace_installer.enable_files_in_repos()
         workspace_installer._config = config  # type: ignore[has-type]
         workspace_installer._write_config(overwrite=False)
         workspace_installer.current_config.override_clusters = override_clusters
@@ -424,7 +424,7 @@ def _name(self, name: str) -> str:
         return f"[{self._prefix.upper()}] {name}"
 
     def _configure(self):
-        self._enable_files_in_repos()
+        self.enable_files_in_repos()
         ws_file_url = self.notebook_link(self.config_file)
         try:
             if "version: 1" in self._raw_previous_config():
@@ -1041,7 +1041,7 @@ def validate_and_run(self, step: str):
             self.run_workflow(step)
 
     @staticmethod
-    def _enable_files_in_repos():
+    def enable_files_in_repos():
         client = ApiClient()
 
         # check if "enableProjectTypeInWorkspace" and "enableWorkspaceFilesystem" are set to false

tests/integration/test_installation.py

@@ -6,7 +6,6 @@
 
 import pytest
 from databricks.labs.blueprint.parallel import Threads
-from databricks.sdk.core import ApiClient
 from databricks.sdk.errors import InvalidParameterValue, NotFound
 from databricks.sdk.retries import retried
 from databricks.sdk.service.iam import PermissionLevel
@@ -183,32 +182,10 @@ def test_uninstallation(ws, sql_backend, new_installation):
         sql_backend.execute(f"show tables from hive_metastore.{install.current_config.inventory_database}")
 
 
-def test_files_in_repos_enablement(ws, sql_backend, new_installation):
-    client = ApiClient()
-
-    project_type = client.do("GET", "/api/2.0/workspace-conf", {"keys": "enableProjectTypeInWorkspace"})
-    workspace_file_system = client.do("GET", "/api/2.0/workspace-conf", {"keys": "enableWorkspaceFilesystem"})
-
-    if project_type["enableProjectTypeInWorkspace"] != "false":
-        client.do(
-            "PATCH",
-            "/api/2.0/workspace-conf",
-            body={"enableProjectTypeInWorkspace": "false"},
-            headers={"Content-Type": "application/json"},
-        )
-
-    if workspace_file_system["enableWorkspaceFilesystem"] != "false":
-        client.do(
-            "PATCH",
-            "/api/2.0/workspace-conf",
-            body={"enableWorkspaceFilesystem": "false"},
-            headers={"Content-Type": "application/json"},
-        )
-
+def test_files_in_repos_enablement(new_installation):
     try:
         install = new_installation()
         install.uninstall()
     except NotFound as e:
         raise AssertionError() from e
-
     assert True

tests/unit/test_install.py

@@ -1,4 +1,5 @@
 import io
+import random
 from datetime import timedelta
 from pathlib import Path
 from typing import Any
@@ -76,10 +77,26 @@ def mock_clusters():
     ]
 
 
+def get_workspace_config(*args, **kwargs):
+    if args[0] == "GET":
+        if args[2]["keys"] == "enableProjectTypeInWorkspace":
+            random_return_value = random.choice(["true", "false"])
+            return {"enableProjectTypeInWorkspace": random_return_value}
+        elif args[2]["keys"] == "enableWorkspaceFilesystem":
+            random_return_value = random.choice(["true", "false"])
+            return {"enableWorkspaceFilesystem": random_return_value}
+    elif args[0] == "PATCH":
+        if "enableProjectTypeInWorkspace" in kwargs["body"]:
+            return {"enableProjectTypeInWorkspace": kwargs["body"]["enableProjectTypeInWorkspace"]}
+        elif "enableWorkspaceFilesystem" in kwargs["body"]:
+            return {"enableWorkspaceFilesystem": kwargs["body"]["enableWorkspaceFilesystem"]}
+
+
 @pytest.fixture
 def ws(mocker):
     ws = mocker.patch("databricks.sdk.WorkspaceClient.__init__")
-
+    mocker.patch("databricks.sdk.core.ApiClient.__init__", return_value=None)
+    mocker.patch("databricks.sdk.core.ApiClient.do", side_effect=get_workspace_config)
     ws.current_user.me = lambda: iam.User(user_name="me@example.com", groups=[iam.ComplexValue(display="admins")])
     ws.config.host = "https://foo"
     ws.config.is_aws = True
@@ -96,6 +113,7 @@ def ws(mocker):
     ws.query_visualizations.create.return_value = Visualization(id="abc")
     ws.dashboard_widgets.create.return_value = Widget(id="abc")
     ws.clusters.list.return_value = mock_clusters()
+
     return ws
 
 
@@ -1200,3 +1218,10 @@ def test_repair_run_result_state(ws, caplog):
     ws.jobs.list_runs.repair_run = None
     install.repair_run("assessment")
     assert "Please try after sometime" in caplog.text
+
+
+def test_enable_files_in_repos(ws, mocker):
+    mocker.patch("databricks.sdk.core.ApiClient.do", side_effect=get_workspace_config)
+    install = WorkspaceInstaller(ws, verify_timeout=timedelta(seconds=5))
+    install.enable_files_in_repos()
+    assert True