diff --git a/autoconf/setup_colab.py b/autoconf/setup_colab.py index 5cb8c2a..3814550 100644 --- a/autoconf/setup_colab.py +++ b/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/``); 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, diff --git a/test_autoconf/test_setup_colab.py b/test_autoconf/test_setup_colab.py index ba7b7e7..6194c84 100644 --- a/test_autoconf/test_setup_colab.py +++ b/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" + )