From bf1eb22fab08deebf6da5a96745f407140869e80 Mon Sep 17 00:00:00 2001
From: Douglas Moore <dmoore247@users.noreply.github.com>
Date: Tue, 3 Oct 2023 19:27:11 -0400
Subject: [PATCH] Improved test coverage for installer (#371)

add test for `_configure_inventory_database()`
add test for `run_for_config()`
---
 src/databricks/labs/ucx/install.py | 27 ++++++++++-------
 tests/unit/test_install.py         | 48 ++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 11 deletions(-)

diff --git a/src/databricks/labs/ucx/install.py b/src/databricks/labs/ucx/install.py
index b1886aeadd..5f41ae5f60 100644
--- a/src/databricks/labs/ucx/install.py
+++ b/src/databricks/labs/ucx/install.py
@@ -167,6 +167,21 @@ def _current_config(self):
     def _name(self, name: str) -> str:
         return f"[{self._prefix.upper()}][{self._short_name}] {name}"
 
+    def _configure_inventory_database(self):
+        counter = 0
+        inventory_database = None
+        while True:
+            inventory_database = self._question("Inventory Database stored in hive_metastore", default="ucx")
+            if re.match(r"^\w+$", inventory_database):
+                break
+            else:
+                print(f"{inventory_database} is not a valid database name")
+                counter = counter + 1
+                if counter > NUM_USER_ATTEMPTS:
+                    msg = "Exceeded max tries to get a valid database name, try again later."
+                    raise SystemExit(msg)
+        return inventory_database
+
     def _configure(self):
         ws_file_url = self._notebook_link(self._config_file)
         try:
@@ -180,17 +195,7 @@ def _configure(self):
                 raise err
 
         logger.info("Please answer a couple of questions to configure Unity Catalog migration")
-        counter = 0
-        while True:
-            inventory_database = self._question("Inventory Database stored in hive_metastore", default="ucx")
-            if re.match(r"^\w+$", inventory_database):
-                break
-            else:
-                print(f"{inventory_database} is not a valid database name")
-                counter = counter + 1
-                if counter > NUM_USER_ATTEMPTS:
-                    msg = "Exceeded max tries to get a valid database name, try again later."
-                    raise SystemExit(msg)
+        inventory_database = self._configure_inventory_database()
 
         pro_warehouses = {"[Create new PRO SQL warehouse]": "create_new"} | {
             f"{_.name} ({_.id}, {_.warehouse_type.value}, {_.state.value})": _.id
diff --git a/tests/unit/test_install.py b/tests/unit/test_install.py
index 8907603d45..c98ffba9bc 100644
--- a/tests/unit/test_install.py
+++ b/tests/unit/test_install.py
@@ -23,6 +23,54 @@
 from databricks.labs.ucx.install import WorkspaceInstaller
 
 
+def mock_ws(mocker):
+    ws = mocker.Mock()
+    ws = mocker.patch("databricks.sdk.WorkspaceClient.__init__")
+
+    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
+    config_bytes = yaml.dump(WorkspaceConfig(inventory_database="a", groups=GroupsConfig(auto=True)).as_dict()).encode(
+        "utf8"
+    )
+    ws.workspace.download = lambda _: io.BytesIO(config_bytes)
+    ws.workspace.get_status = lambda _: ObjectInfo(object_id=123)
+    ws.data_sources.list = lambda: [DataSource(id="bcd", warehouse_id="abc")]
+    ws.warehouses.list = lambda **_: [EndpointInfo(id="abc", warehouse_type=EndpointInfoWarehouseType.PRO)]
+    ws.dashboards.create.return_value = Dashboard(id="abc")
+    ws.queries.create.return_value = Query(id="abc")
+    ws.query_visualizations.create.return_value = Visualization(id="abc")
+    ws.dashboard_widgets.create.return_value = Widget(id="abc")
+    return ws
+
+
+def test_run_for_config(mocker, tmp_path):
+    # run_for_config(ws: WorkspaceClient, config: WorkspaceConfig, *, prefix="ucx") -> "WorkspaceInstaller":
+    ws = mock_ws(mocker)
+
+    install = WorkspaceInstaller(ws)
+    wc = WorkspaceConfig(inventory_database="a", groups=GroupsConfig(auto=True))
+    return_value = install.run_for_config(ws, wc)
+    assert return_value
+
+
+def test_install_database_happy(mocker, tmp_path):
+    ws = mocker.Mock()
+    install = WorkspaceInstaller(ws)
+    mocker.patch("builtins.input", return_value="ucx")
+    res = install._configure_inventory_database()
+    assert "ucx" == res
+
+
+def test_install_database_unhappy(mocker, tmp_path):
+    ws = mocker.Mock()
+    install = WorkspaceInstaller(ws)
+    mocker.patch("builtins.input", return_value="main.ucx")
+
+    with pytest.raises(SystemExit):
+        install._configure_inventory_database()
+
+
 def test_build_wheel(mocker, tmp_path):
     ws = mocker.Mock()
     install = WorkspaceInstaller(ws)