andrew/get chrome path update

CHANGELOG.txt

@@ -1,3 +1,5 @@
+- Change chrome download path to use XDG cache dir
+- Don't download chrome if we already have that version: add force argument
 v1.2.1
 - Use custom threadpool for functions that could be running during shutdown:
   Python's stdlib threadpool isn't available during interpreter shutdown, nor

pyproject.toml

@@ -30,6 +30,7 @@ maintainers = [
 ]
 dependencies = [
   "logistro>=2.0.1",
+  "platformdirs>=4.3.6",
   "simplejson>=3.19.3",
 ]
 

src/choreographer/cli/_cli_utils.py

@@ -23,7 +23,7 @@
 supported_platform_strings = ["linux64", "win32", "win64", "mac-x64", "mac-arm64"]
 
 
-def get_google_supported_platform_string() -> tuple[str, str, str, str]:
+def get_google_supported_platform_string() -> str | None:
     arch_size_detected = "64" if sys.maxsize > 2**32 else "32"
     arch_detected = "arm" if platform.processor() == "arm" else "x"
 
@@ -39,16 +39,17 @@ def get_google_supported_platform_string() -> tuple[str, str, str, str]:
     if chrome_platform_detected in supported_platform_strings:
         platform_string = chrome_platform_detected
 
-    return platform_string, arch_size_detected, platform.processor(), platform.system()
+    return platform_string
 
 
 def get_chrome_download_path() -> Path | None:
-    _chrome_platform_detected, _, _, _ = get_google_supported_platform_string()
+    _chrome_platform_detected = get_google_supported_platform_string()
 
     if not _chrome_platform_detected:
         return None
 
-    _default_exe_path = Path()
+    _default_exe_path = default_download_path
+    _default_exe_path.mkdir(parents=True, exist_ok=True)
 
     if platform.system().startswith("Linux"):
         _default_exe_path = (
@@ -85,26 +86,26 @@ def _extract_member(self, member, targetpath, pwd):  # type: ignore [no-untyped-
         return path
 
 
-def get_chrome_sync(  # noqa: PLR0912, C901
+def get_chrome_sync(  # noqa: C901, PLR0912
     arch: str | None = None,
     i: int | None = None,
     path: str | Path = default_download_path,
     *,
     verbose: bool = False,
+    force: bool = False,
 ) -> Path | str:
     """Download chrome synchronously: see `get_chrome()`."""
-    if not arch:
-        arch, _, _, _ = get_google_supported_platform_string()
+    if isinstance(path, str):
+        path = Path(path)
 
+    arch = arch or get_google_supported_platform_string()
     if not arch:
         raise RuntimeError(
             "You must specify an arch, one of: "
             f"{', '.join(supported_platform_strings)}. "
             f"Detected {arch} is not supported.",
         )
 
-    if isinstance(path, str):
-        path = Path(path)
     if i:
         _logger.info("Loading chrome from list")
         browser_list = json.loads(
@@ -115,17 +116,16 @@ def get_chrome_sync(  # noqa: PLR0912, C901
         version_obj = browser_list["versions"][i]
     else:
         _logger.info("Using last known good version of chrome")
-        with (
-            Path(__file__).resolve().parent.parent
-            / "resources"
-            / "last_known_good_chrome.json"
-        ).open() as f:
-            version_obj = json.load(f)
-    if verbose:
-        print(version_obj["version"])  # noqa: T201 allow print in cli
-        print(version_obj["revision"])  # noqa: T201 allow print in cli
+        version_obj = json.loads(
+            (
+                Path(__file__).resolve().parent.parent
+                / "resources"
+                / "last_known_good_chrome.json"
+            ).read_text(),
+        )
+    version_string = f"{version_obj['version']}\n{version_obj['revision']}"
     chromium_sources = version_obj["downloads"]["chrome"]
-    url = ""
+
     for src in chromium_sources:
         if src["platform"] == arch:
             url = src["url"]
@@ -137,19 +137,16 @@ def get_chrome_sync(  # noqa: PLR0912, C901
             f"{arch} is not supported.",
         )
 
-    if not path.exists():
-        path.mkdir(parents=True)
-    filename = path / "chrome.zip"
-    with urllib.request.urlopen(url) as response, filename.open("wb") as out_file:  # noqa: S310 audit url
-        shutil.copyfileobj(response, out_file)
-    with _ZipFilePermissions(filename, "r") as zip_ref:
-        zip_ref.extractall(path)
-    filename.unlink()
+    if verbose:
+        print(version_string)  # noqa: T201 allow print in cli
+    version_tag = path / "version_tag.txt"
+
+    path.mkdir(parents=True, exist_ok=True)
 
     if arch.startswith("linux"):
-        exe_name = path / f"chrome-{arch}" / "chrome"
+        exe_path = path / f"chrome-{arch}" / "chrome"
     elif arch.startswith("mac"):
-        exe_name = (
+        exe_path = (
             path
             / f"chrome-{arch}"
             / "Google Chrome for Testing.app"
@@ -158,10 +155,37 @@ def get_chrome_sync(  # noqa: PLR0912, C901
             / "Google Chrome for Testing"
         )
     elif arch.startswith("win"):
-        exe_name = path / f"chrome-{arch}" / "chrome.exe"
+        exe_path = path / f"chrome-{arch}" / "chrome.exe"
     else:
         raise RuntimeError("Couldn't calculate exe_name, unsupported architecture.")
-    return exe_name
+
+    if (
+        exe_path.exists()
+        and version_tag.is_file()
+        and version_tag.read_text() == version_string
+        and not force
+    ):
+        return exe_path
+    else:
+        if exe_path.exists():  # delete it
+            if exe_path.is_dir():
+                shutil.rmtree(exe_path)
+            else:
+                exe_path.unlink()
+            # It really should always be a dir but in testing we fake it
+        if version_tag.exists():  # delete it
+            version_tag.unlink()
+
+    # Download
+    zip_path = path / "chrome.zip"
+    with urllib.request.urlopen(url) as response, zip_path.open("wb") as out_file:  # noqa: S310 audit url
+        shutil.copyfileobj(response, out_file)
+    with _ZipFilePermissions(zip_path, "r") as zip_ref:
+        zip_ref.extractall(path)
+    zip_path.unlink()
+    version_tag.write_text(version_string)
+
+    return exe_path
 
 
 async def get_chrome(
@@ -170,6 +194,7 @@ async def get_chrome(
     path: str | Path = default_download_path,
     *,
     verbose: bool = False,
+    force: bool = False,
 ) -> Path | str:
     """
     Download google chrome from google-chrome-for-testing server.
@@ -180,10 +205,18 @@ async def get_chrome(
            still in the testing directory.
         path: where to download it too (the folder).
         verbose: print out version found
+        force: download chrome again even if already present at that version
 
     """
     loop = asyncio.get_running_loop()
-    fn = partial(get_chrome_sync, arch=arch, i=i, path=path, verbose=verbose)
+    fn = partial(
+        get_chrome_sync,
+        arch=arch,
+        i=i,
+        path=path,
+        verbose=verbose,
+        force=force,
+    )
     return await loop.run_in_executor(
         executor=None,
         func=fn,
@@ -231,12 +264,21 @@ def get_chrome_cli() -> None:
         action="store_true",
         help="Display found version number if using -i (to stdout)",
     )
+    parser.add_argument(
+        "-f",
+        "--force",
+        dest="force",
+        action="store_true",
+        default=False,
+        help="Force download even if already present.",
+    )
     parser.set_defaults(path=default_download_path)
     parser.set_defaults(arch=None)
     parser.set_defaults(verbose=False)
     parsed = parser.parse_args()
     i = parsed.i
     arch = parsed.arch
     path = Path(parsed.path)
+    force = parsed.force
     verbose = parsed.verbose
-    print(get_chrome_sync(arch=arch, i=i, path=path, verbose=verbose))  # noqa: T201 allow print in cli
+    print(get_chrome_sync(arch=arch, i=i, path=path, verbose=verbose, force=force))  # noqa: T201 allow print in cli

src/choreographer/cli/defaults.py

@@ -2,5 +2,12 @@
 
 from pathlib import Path
 
-default_download_path = Path(__file__).resolve().parent / "browser_exe"
+from platformdirs import PlatformDirs
+
+default_download_path = (
+    Path(
+        PlatformDirs("choreographer", "plotly").user_data_dir,
+    )
+    / "deps"
+)
 """The path where we download chrome if no path argument is supplied."""

tests/test_get_chrome.py

@@ -0,0 +1,78 @@
+import io
+import zipfile
+from unittest.mock import patch
+
+import pytest
+
+from choreographer.cli._cli_utils import get_chrome_sync
+
+
+@pytest.fixture
+def mock_last_known_good_json():
+    """Mock the last known good chrome version."""
+    return {
+        "version": "135.0.7011.0",
+        "revision": "1418433",
+        "downloads": {
+            "chrome": [
+                {"platform": "linux64", "url": "https://example.com/linux64.zip"},
+                {"platform": "mac-arm64", "url": "https://example.com/mac-arm64.zip"},
+                {"platform": "win64", "url": "https://example.com/win64.zip"},
+            ],
+        },
+    }
+
+
+# Thanks Claude!
+def test_get_chrome_sync_download_behavior(
+    tmp_path,
+    mock_last_known_good_json,
+):
+    """Test chrome download/skip behavior: existing, force, version change."""
+    exe_path = tmp_path / "chrome-linux64" / "chrome"
+    version_tag = tmp_path / "version_tag.txt"
+
+    # Setup: create a fake existing installation with matching version
+    exe_path.parent.mkdir(parents=True)
+    exe_path.write_text("old chrome")
+    version_tag.write_text("135.0.7011.0\n1418433")
+
+    # Mock the URL download
+    def create_mock_zip_response():
+        # Create a fresh zip buffer each time
+        zip_buffer = io.BytesIO()
+        with zipfile.ZipFile(zip_buffer, "w") as zf:
+            zf.writestr("chrome-linux64/chrome", "newly downloaded chrome")
+        zip_buffer.seek(0)
+        # Return BytesIO directly - it already has all file-like methods needed
+        return zip_buffer
+
+    # Patch json.load to return our mock data (avoid broad Path.read_text patch)
+    with patch("json.loads", return_value=mock_last_known_good_json), patch(
+        "urllib.request.urlopen",
+        side_effect=lambda url: create_mock_zip_response(),  # noqa: ARG005
+    ):
+        # a) First call without force - should return existing, no download
+        result = get_chrome_sync(arch="linux64", path=tmp_path, force=False)
+        assert result == exe_path
+        assert exe_path.read_text() == "old chrome"  # Should still be old
+
+        # b) Call with force=True - should download and replace
+        result = get_chrome_sync(arch="linux64", path=tmp_path, force=True)
+        assert result == exe_path
+        assert exe_path.read_text() == "newly downloaded chrome"
+        assert version_tag.read_text() == "135.0.7011.0\n1418433"
+
+        # c) Call again without force - should return existing, no re-download
+        # (Modify the file to verify it doesn't get re-downloaded)
+        exe_path.write_text("manually modified chrome")
+        result = get_chrome_sync(arch="linux64", path=tmp_path, force=False)
+        assert result == exe_path
+        assert exe_path.read_text() == "manually modified chrome"
+
+        # d) Change version on disk - should trigger download (version mismatch)
+        version_tag.write_text("999.0.0.0\n999999")  # Old/different version
+        result = get_chrome_sync(arch="linux64", path=tmp_path, force=False)
+        assert result == exe_path
+        assert exe_path.read_text() == "newly downloaded chrome"
+        assert version_tag.read_text() == "135.0.7011.0\n1418433"  # Updated to new