feat: workspace_dir override in setup_colab.setup (CI Colab simulation)

autoconf/setup_colab.py

@@ -266,7 +266,11 @@ def _colab_setup(
     )
 
 
-def setup(project: str, raise_error_if_not_gpu: bool = False) -> None:
+def setup(
+    project: str,
+    raise_error_if_not_gpu: bool = False,
+    workspace_dir: str = None,
+) -> None:
     """
     Set up Google Colab for a PyAuto notebook repository.
 
@@ -278,6 +282,11 @@ def setup(project: str, raise_error_if_not_gpu: bool = False) -> None:
     raise_error_if_not_gpu
         If True, raise instead of continuing when Colab has no GPU / TPU
         runtime enabled.
+    workspace_dir
+        Where to clone the workspace. Defaults to the project's Colab
+        directory (``/content/<workspace>``); overridden by CI environments
+        that simulate the Colab bootstrap (e.g. PyAutoHeart's verify_install
+        check F), where ``/content`` is not writable.
     """
     try:
         spec = _PROJECTS[project]
@@ -289,7 +298,7 @@ def setup(project: str, raise_error_if_not_gpu: bool = False) -> None:
     _colab_setup(
         project_name=spec["project_name"],
         workspace_repo=spec["workspace_repo"],
-        workspace_dir=spec["workspace_dir"],
+        workspace_dir=workspace_dir or spec["workspace_dir"],
         packages=spec["packages"],
         top_package=spec["top_package"],
         raise_error_if_not_gpu=raise_error_if_not_gpu,

test_autoconf/test_setup_colab.py

@@ -159,3 +159,18 @@ def test_falls_back_when_version_unknown(self, tmp_path):
                 setup_colab._clone_workspace("repo_url", target, "autolens")
         run.assert_called_once()
         assert "--branch" not in run.call_args[0][0]
+
+
+class TestWorkspaceDirOverride:
+    def test_override_threads_to_colab_setup(self):
+        with mock.patch.object(setup_colab, "_colab_setup") as colab_setup:
+            setup_colab.setup("autolens", workspace_dir="/tmp/sim_ws")
+        assert colab_setup.call_args[1]["workspace_dir"] == "/tmp/sim_ws"
+
+    def test_default_is_the_registry_colab_dir(self):
+        with mock.patch.object(setup_colab, "_colab_setup") as colab_setup:
+            setup_colab.setup("autolens")
+        assert (
+            colab_setup.call_args[1]["workspace_dir"]
+            == "/content/autolens_workspace"
+        )