From 30ca9a724335ea929d965e0fb4b70a3682cfa64f Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 9 Jul 2026 17:08:06 +0100 Subject: [PATCH] Add workspace_dir override to setup_colab.setup for CI Colab simulation The registry hardcodes Colab's /content/ directories, which CI cannot write; PyAutoHeart's verify_install check F simulates the Colab bootstrap in a throwaway venv and passes its own clone target through this parameter. Default behavior unchanged. Part of PyAutoLabs/PyAutoHeart#45 Co-Authored-By: Claude Fable 5 --- autoconf/setup_colab.py | 13 +++++++++++-- test_autoconf/test_setup_colab.py | 15 +++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) 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" + )