feature: cache downloaded wheel information (#2276)

Co-authored-by: Bernát Gábor <gaborjbernat@gmail.com>

2 files changed, 160 insertions, 61 deletions

diff --git a/tests/unit/seed/wheels/test_acquire.py b/tests/unit/seed/wheels/test_acquire.py
index e255311..dcce0f6 100644
--- a/tests/unit/seed/wheels/test_acquire.py
+++ b/tests/unit/seed/wheels/test_acquire.py
@@ -2,16 +2,25 @@ from __future__ import absolute_import, unicode_literals
 
 import os
 import sys
+from datetime import datetime
 from subprocess import CalledProcessError
 
 import pytest
 
+from virtualenv.app_data import AppDataDiskFolder
+from virtualenv.info import IS_PYPY, PY2
 from virtualenv.seed.wheels.acquire import download_wheel, get_wheel, pip_wheel_env_run
 from virtualenv.seed.wheels.embed import BUNDLE_FOLDER, get_embed_wheel
+from virtualenv.seed.wheels.periodic_update import dump_datetime
 from virtualenv.seed.wheels.util import Wheel, discover_wheels
 from virtualenv.util.path import Path
 
 
+@pytest.fixture(autouse=True)
+def fake_release_date(mocker):
+    mocker.patch("virtualenv.seed.wheels.periodic_update.release_date_for_wheel_path", return_value=None)
+
+
 def test_pip_wheel_env_run_could_not_find(session_app_data, mocker):
     mocker.patch("virtualenv.seed.wheels.acquire.from_bundle", return_value=None)
     with pytest.raises(RuntimeError, match="could not find the embedded pip"):
@@ -74,24 +83,64 @@ def test_download_fails(mocker, for_py_version, session_app_data):
 @pytest.fixture
 def downloaded_wheel(mocker):
     wheel = Wheel.from_path(Path("setuptools-0.0.0-py2.py3-none-any.whl"))
-    mocker.patch("virtualenv.seed.wheels.acquire.download_wheel", return_value=wheel)
-    yield wheel
+    yield wheel, mocker.patch("virtualenv.seed.wheels.acquire.download_wheel", return_value=wheel)
 
 
 @pytest.mark.parametrize("version", ["bundle", "0.0.0"])
-def test_get_wheel_download_called(for_py_version, session_app_data, downloaded_wheel, version):
+def test_get_wheel_download_called(mocker, for_py_version, session_app_data, downloaded_wheel, version):
     distribution = "setuptools"
+    write = mocker.patch("virtualenv.app_data.via_disk_folder.JSONStoreDisk.write")
     wheel = get_wheel(distribution, version, for_py_version, [], True, session_app_data, False, os.environ)
     assert wheel is not None
-    assert wheel.name == downloaded_wheel.name
+    assert wheel.name == downloaded_wheel[0].name
+    assert downloaded_wheel[1].call_count == 1
+    assert write.call_count == 1
 
 
 @pytest.mark.parametrize("version", ["embed", "pinned"])
-def test_get_wheel_download_not_called(for_py_version, session_app_data, downloaded_wheel, version):
+def test_get_wheel_download_not_called(mocker, for_py_version, session_app_data, downloaded_wheel, version):
     distribution = "setuptools"
     expected = get_embed_wheel(distribution, for_py_version)
     if version == "pinned":
         version = expected.version
+    write = mocker.patch("virtualenv.app_data.via_disk_folder.JSONStoreDisk.write")
     wheel = get_wheel(distribution, version, for_py_version, [], True, session_app_data, False, os.environ)
     assert wheel is not None
     assert wheel.name == expected.name
+    assert downloaded_wheel[1].call_count == 0
+    assert write.call_count == 0
+
+
+@pytest.mark.skipif(IS_PYPY and PY2, reason="mocker.spy failing on PyPy 2.x")
+def test_get_wheel_download_cached(tmp_path, freezer, mocker, for_py_version, downloaded_wheel):
+    from virtualenv.app_data.via_disk_folder import JSONStoreDisk
+
+    app_data = AppDataDiskFolder(folder=str(tmp_path))
+    expected = downloaded_wheel[0]
+    write = mocker.spy(JSONStoreDisk, "write")
+    # 1st call, not cached, download is called
+    wheel = get_wheel(expected.distribution, expected.version, for_py_version, [], True, app_data, False, os.environ)
+    assert wheel is not None
+    assert wheel.name == expected.name
+    assert downloaded_wheel[1].call_count == 1
+    assert write.call_count == 1
+    # 2nd call, cached, download is not called
+    wheel = get_wheel(expected.distribution, expected.version, for_py_version, [], True, app_data, False, os.environ)
+    assert wheel is not None
+    assert wheel.name == expected.name
+    assert downloaded_wheel[1].call_count == 1
+    assert write.call_count == 1
+    wrote_json = write.call_args[0][1]
+    assert wrote_json == {
+        "completed": None,
+        "periodic": None,
+        "started": None,
+        "versions": [
+            {
+                "filename": expected.name,
+                "release_date": None,
+                "found_date": dump_datetime(datetime.now()),
+                "source": "download",
+            },
+        ],
+    }
diff --git a/tests/unit/seed/wheels/test_periodic_update.py b/tests/unit/seed/wheels/test_periodic_update.py
index e35dfed..7adb359 100644
--- a/tests/unit/seed/wheels/test_periodic_update.py
+++ b/tests/unit/seed/wheels/test_periodic_update.py
@@ -170,9 +170,9 @@ def test_manual_update_honored(mocker, session_app_data, for_py_version):
     assert str(result.path) == expected_path
 
 
-def wheel_path(wheel, of):
+def wheel_path(wheel, of, pre_release=""):
     new_version = ".".join(str(i) for i in (tuple(sum(x) for x in zip_longest(wheel.version_tuple, of, fillvalue=0))))
-    new_name = wheel.name.replace(wheel.version, new_version)
+    new_name = wheel.name.replace(wheel.version, new_version + pre_release)
     return str(wheel.path.parent / new_name)
 
 
@@ -488,21 +488,25 @@ def test_get_release_fails(mocker, caplog):
     assert repr(exc) in caplog.text
 
 
+def mock_download(mocker, pip_version_remote):
+    def download():
+        index = 0
+        while True:
+            path = pip_version_remote[index]
+            index += 1
+            yield Wheel(Path(path))
+
+    do = download()
+    return mocker.patch("virtualenv.seed.wheels.acquire.download_wheel", side_effect=lambda *a, **k: next(do))
+
+
 def test_download_stop_with_embed(tmp_path, mocker, freezer):
     freezer.move_to(_UP_NOW)
     wheel = get_embed_wheel("pip", "3.9")
     app_data_outer = AppDataDiskFolder(str(tmp_path / "app"))
     pip_version_remote = [wheel_path(wheel, (0, 0, 2)), wheel_path(wheel, (0, 0, 1)), wheel_path(wheel, (-1, 0, 0))]
-    at = {"index": 0}
-
-    def download():
-        while True:
-            path = pip_version_remote[at["index"]]
-            at["index"] += 1
-            yield Wheel(Path(path))
 
-    do = download()
-    download_wheel = mocker.patch("virtualenv.seed.wheels.acquire.download_wheel", side_effect=lambda *a, **k: next(do))
+    download_wheel = mock_download(mocker, pip_version_remote)
     url_o = mocker.patch("virtualenv.seed.wheels.periodic_update.urlopen", side_effect=URLError("unavailable"))
 
     last_update = _UP_NOW - timedelta(days=14)
@@ -519,32 +523,71 @@ def test_download_stop_with_embed(tmp_path, mocker, freezer):
     assert write.call_count == 1
 
 
+def test_download_manual_stop_after_one_download(tmp_path, mocker, freezer):
+    freezer.move_to(_UP_NOW)
+    wheel = get_embed_wheel("pip", "3.9")
+    app_data_outer = AppDataDiskFolder(str(tmp_path / "app"))
+    pip_version_remote = [wheel_path(wheel, (0, 1, 1))]
+
+    download_wheel = mock_download(mocker, pip_version_remote)
+    url_o = mocker.patch("virtualenv.seed.wheels.periodic_update.urlopen", side_effect=URLError("unavailable"))
+
+    last_update = _UP_NOW - timedelta(days=14)
+    u_log = UpdateLog(started=last_update, completed=last_update, versions=[], periodic=True)
+    read_dict = mocker.patch("virtualenv.app_data.via_disk_folder.JSONStoreDisk.read", return_value=u_log.to_dict())
+    write = mocker.patch("virtualenv.app_data.via_disk_folder.JSONStoreDisk.write")
+
+    do_update("pip", "3.9", str(wheel.path), str(app_data_outer), [], False)
+
+    assert download_wheel.call_count == 1
+    assert url_o.call_count == 2
+    assert read_dict.call_count == 1
+    assert write.call_count == 1
+
+
+def test_download_manual_ignores_pre_release(tmp_path, mocker, freezer):
+    freezer.move_to(_UP_NOW)
+    wheel = get_embed_wheel("pip", "3.9")
+    app_data_outer = AppDataDiskFolder(str(tmp_path / "app"))
+    pip_version_remote = [wheel_path(wheel, (0, 0, 1))]
+    pip_version_pre = NewVersion(Path(wheel_path(wheel, (0, 1, 0), "b1")).name, _UP_NOW, None, "downloaded")
+
+    download_wheel = mock_download(mocker, pip_version_remote)
+    url_o = mocker.patch("virtualenv.seed.wheels.periodic_update.urlopen", side_effect=URLError("unavailable"))
+
+    last_update = _UP_NOW - timedelta(days=14)
+    u_log = UpdateLog(started=last_update, completed=last_update, versions=[pip_version_pre], periodic=True)
+    read_dict = mocker.patch("virtualenv.app_data.via_disk_folder.JSONStoreDisk.read", return_value=u_log.to_dict())
+    write = mocker.patch("virtualenv.app_data.via_disk_folder.JSONStoreDisk.write")
+
+    do_update("pip", "3.9", str(wheel.path), str(app_data_outer), [], False)
+
+    assert download_wheel.call_count == 1
+    assert url_o.call_count == 2
+    assert read_dict.call_count == 1
+    assert write.call_count == 1
+    wrote_json = write.call_args[0][0]
+    assert wrote_json["versions"] == [
+        {
+            "filename": Path(pip_version_remote[0]).name,
+            "release_date": None,
+            "found_date": dump_datetime(_UP_NOW),
+            "source": "manual",
+        },
+        pip_version_pre.to_dict(),
+    ]
+
+
 def test_download_periodic_stop_at_first_usable(tmp_path, mocker, freezer):
     freezer.move_to(_UP_NOW)
     wheel = get_embed_wheel("pip", "3.9")
     app_data_outer = AppDataDiskFolder(str(tmp_path / "app"))
     pip_version_remote = [wheel_path(wheel, (0, 1, 1)), wheel_path(wheel, (0, 1, 0))]
     rel_date_remote = [_UP_NOW - timedelta(days=1), _UP_NOW - timedelta(days=30)]
-    at = {"download": 0, "release_date": 0}
-
-    def download():
-        while True:
-            path = pip_version_remote[at["download"]]
-            at["download"] += 1
-            yield Wheel(Path(path))
 
-    download_gen = download()
-    download_wheel = mocker.patch(
-        "virtualenv.seed.wheels.acquire.download_wheel", side_effect=lambda *a, **k: next(download_gen)
-    )
-
-    def rel_date():
-        while True:
-            value = rel_date_remote[at["release_date"]]
-            at["release_date"] += 1
-            yield value
+    download_wheel = mock_download(mocker, pip_version_remote)
 
-    rel_date_gen = rel_date()
+    rel_date_gen = iter(rel_date_remote)
     release_date = mocker.patch(
         "virtualenv.seed.wheels.periodic_update.release_date_for_wheel_path",
         side_effect=lambda *a, **k: next(rel_date_gen),
@@ -559,51 +602,58 @@ def test_download_periodic_stop_at_first_usable(tmp_path, mocker, freezer):
 
     assert download_wheel.call_count == 2
     assert release_date.call_count == 2
-
     assert read_dict.call_count == 1
     assert write.call_count == 1
 
 
-def test_download_manual_stop_at_first_usable(tmp_path, mocker, freezer):
+def test_download_periodic_stop_at_first_usable_with_previous_minor(tmp_path, mocker, freezer):
     freezer.move_to(_UP_NOW)
     wheel = get_embed_wheel("pip", "3.9")
     app_data_outer = AppDataDiskFolder(str(tmp_path / "app"))
-    pip_version_remote = [wheel_path(wheel, (0, 1, 1))]
-    rel_date_remote = [_UP_NOW + timedelta(hours=1)]
-    at = {"download": 0, "release_date": 0}
-
-    def download():
-        while True:
-            path = pip_version_remote[at["download"]]
-            at["download"] += 1
-            yield Wheel(Path(path))
-
-    download_gen = download()
-    download_wheel = mocker.patch(
-        "virtualenv.seed.wheels.acquire.download_wheel", side_effect=lambda *a, **k: next(download_gen)
-    )
+    pip_version_remote = [wheel_path(wheel, (0, 1, 1)), wheel_path(wheel, (0, 1, 0)), wheel_path(wheel, (0, -1, 0))]
+    rel_date_remote = [_UP_NOW - timedelta(days=1), _UP_NOW - timedelta(days=30), _UP_NOW - timedelta(days=40)]
+    downloaded_versions = [
+        NewVersion(Path(pip_version_remote[2]).name, rel_date_remote[2], None, "download"),
+        NewVersion(Path(pip_version_remote[0]).name, rel_date_remote[0], None, "download"),
+    ]
 
-    def rel_date():
-        while True:
-            value = rel_date_remote[at["release_date"]]
-            at["release_date"] += 1
-            yield value
+    download_wheel = mock_download(mocker, pip_version_remote)
 
-    rel_date_gen = rel_date()
+    rel_date_gen = iter(rel_date_remote)
     release_date = mocker.patch(
         "virtualenv.seed.wheels.periodic_update.release_date_for_wheel_path",
         side_effect=lambda *a, **k: next(rel_date_gen),
     )
 
     last_update = _UP_NOW - timedelta(days=14)
-    u_log = UpdateLog(started=last_update, completed=last_update, versions=[], periodic=True)
+    u_log = UpdateLog(
+        started=last_update,
+        completed=last_update,
+        versions=downloaded_versions,
+        periodic=True,
+    )
     read_dict = mocker.patch("virtualenv.app_data.via_disk_folder.JSONStoreDisk.read", return_value=u_log.to_dict())
     write = mocker.patch("virtualenv.app_data.via_disk_folder.JSONStoreDisk.write")
 
-    do_update("pip", "3.9", str(wheel.path), str(app_data_outer), [], False)
-
-    assert download_wheel.call_count == 1
-    assert release_date.call_count == 1
+    do_update("pip", "3.9", str(wheel.path), str(app_data_outer), [], True)
 
+    assert download_wheel.call_count == 2
+    assert release_date.call_count == 2
     assert read_dict.call_count == 1
     assert write.call_count == 1
+    wrote_json = write.call_args[0][0]
+    assert wrote_json["versions"] == [
+        {
+            "filename": Path(pip_version_remote[0]).name,
+            "release_date": dump_datetime(rel_date_remote[0]),
+            "found_date": dump_datetime(_UP_NOW),
+            "source": "periodic",
+        },
+        {
+            "filename": Path(pip_version_remote[1]).name,
+            "release_date": dump_datetime(rel_date_remote[1]),
+            "found_date": dump_datetime(_UP_NOW),
+            "source": "periodic",
+        },
+        downloaded_versions[0].to_dict(),
+    ]