Bug 1592696 - avoid re-fetching profiles when possible r=rwood

This patch uses a cache dir at ~/.condprof-cache so
we avoid re-downloading the same file several times.

Differential Revision: https://phabricator.services.mozilla.com/D52180

--HG--
extra : moz-landing-system : lando

testing/condprofile/condprof/client.py

@@ -13,7 +13,7 @@ import shutil
 
 from condprof import check_install  # NOQA
 from condprof import progress
-from condprof.util import check_exists, download_file, TASK_CLUSTER, get_logger
+from condprof.util import download_file, TASK_CLUSTER, get_logger, ArchiveNotFound
 from condprof.changelog import Changelog
 
 
@@ -26,13 +26,21 @@ CHANGELOG_LINK = (
     ROOT_URL + "/v1/task/" + INDEX_PATH + "/" + PUBLIC_DIR + "/changelog.json"
 )
 DIRECT_LINK = "https://taskcluster-artifacts.net/%(task_id)s/0/public/condprof/"
+CONDPROF_CACHE = "~/.condprof-cache"
 
 
 class ProfileNotFoundError(Exception):
     pass
 
 
-def get_profile(target_dir, platform, scenario, customization="default", task_id=None):
+def get_profile(
+    target_dir,
+    platform,
+    scenario,
+    customization="default",
+    task_id=None,
+    download_cache=True,
+):
     """Extract a conditioned profile in the target directory.
 
     If task_id is provided, will grab the profile from that task. when not
@@ -51,14 +59,21 @@ def get_profile(target_dir, platform, scenario, customization="default", task_id
     else:
         url = DIRECT_LINK % params + filename
 
+    if not download_cache:
         download_dir = tempfile.mkdtemp()
+    else:
+        # using a cache dir in the user home dir
+        download_dir = os.path.expanduser(CONDPROF_CACHE)
+        if not os.path.exists(download_dir):
+            os.makedirs(download_dir)
+
     downloaded_archive = os.path.join(download_dir, filename)
     get_logger().msg("Getting %s" % url)
-    exists, __ = check_exists(url)
-    if exists != 200:
-        raise ProfileNotFoundError(exists)
-
+    try:
         archive = download_file(url, target=downloaded_archive)
+    except ArchiveNotFound:
+        raise ProfileNotFoundError(url)
+
     try:
         with tarfile.open(archive, "r:gz") as tar:
             get_logger().msg("Extracting the tarball content in %s" % target_dir)
@@ -76,6 +91,7 @@ def get_profile(target_dir, platform, scenario, customization="default", task_id
     except (OSError, tarfile.ReadError) as e:
         raise ProfileNotFoundError(str(e))
     finally:
+        if not download_cache:
             shutil.rmtree(download_dir)
     get_logger().msg("Success, we have a profile to work with")
     return target_dir
@@ -85,12 +101,13 @@ def read_changelog(platform):
     params = {"platform": platform}
     changelog_url = CHANGELOG_LINK % params
     get_logger().msg("Getting %s" % changelog_url)
-    exists, __ = check_exists(changelog_url)
-    if exists != 200:
-        raise ProfileNotFoundError(exists)
     download_dir = tempfile.mkdtemp()
     downloaded_changelog = os.path.join(download_dir, "changelog.json")
+    try:
         download_file(changelog_url, target=downloaded_changelog)
+    except ArchiveNotFound:
+        shutil.rmtree(download_dir)
+        raise ProfileNotFoundError(changelog_url)
     return Changelog(download_dir)
 
 

testing/condprofile/condprof/metadata.py

@@ -5,7 +5,7 @@
 """
 import os
 import datetime
-import collections
+from collections.abc import MutableMapping
 import json
 
 from condprof.util import LOG
@@ -14,7 +14,7 @@ from condprof.util import LOG
 METADATA_NAME = "condprofile.json"
 
 
-class Metadata(collections.MutableMapping):
+class Metadata(MutableMapping):
     """ dict-like class that holds metadata for a profile.
     """
 

testing/condprofile/condprof/tests/test_client.py

@@ -0,0 +1,66 @@
+import unittest
+import os
+import tempfile
+import shutil
+import responses
+import re
+
+from condprof.client import get_profile
+
+PROFILE = re.compile("https://index.taskcluster.net/.*/.*tgz")
+with open(os.path.join(os.path.dirname(__file__), "profile.tgz"), "rb") as f:
+    PROFILE_DATA = f.read()
+
+
+class TestClient(unittest.TestCase):
+    def setUp(self):
+        self.target = tempfile.mkdtemp()
+        self.download_dir = os.path.expanduser("~/.condprof-cache")
+        if os.path.exists(self.download_dir):
+            shutil.rmtree(self.download_dir)
+
+        responses.add(
+            responses.GET,
+            PROFILE,
+            body=PROFILE_DATA,
+            headers={"content-length": str(len(PROFILE_DATA)), "ETag": "'12345'"},
+            status=200,
+        )
+
+        responses.add(
+            responses.HEAD,
+            PROFILE,
+            body="",
+            headers={"content-length": str(len(PROFILE_DATA)), "ETag": "'12345'"},
+            status=200,
+        )
+
+    def tearDown(self):
+        shutil.rmtree(self.target)
+        shutil.rmtree(self.download_dir)
+
+    @responses.activate
+    def test_cache(self):
+        download_dir = os.path.expanduser("~/.condprof-cache")
+        if os.path.exists(download_dir):
+            num_elmts = len(os.listdir(download_dir))
+        else:
+            num_elmts = 0
+
+        get_profile(self.target, "win64", "cold", "default")
+
+        # grabbing a profile should generate two files
+        self.assertEqual(len(os.listdir(download_dir)), num_elmts + 2)
+
+        # we do two network calls when getting a file, a HEAD and a GET
+        response_calls = len(responses.calls)
+        self.assertEqual(response_calls, 2)
+
+        # and we should reuse them without downloading the file again
+        get_profile(self.target, "win64", "cold", "default")
+
+        # grabbing a profile should not download new stuff
+        self.assertEqual(len(os.listdir(download_dir)), num_elmts + 2)
+
+        # and do a single extra HEAD call
+        self.assertEqual(len(responses.calls), response_calls + 1)

testing/condprofile/condprof/tests/test_runner.py

@@ -31,11 +31,7 @@ class TestRunner(unittest.TestCase):
         self.archive_dir = tempfile.mkdtemp()
         responses.add(responses.GET, CHANGELOG, json={"error": "not found"}, status=404)
         responses.add(
-            responses.GET,
-            FTP,
-            content_type="application/text/html",
-            body=FTP_PAGE,
-            status=200,
+            responses.GET, FTP, content_type="text/html", body=FTP_PAGE, status=200
         )
 
         responses.add(

testing/condprofile/setup.py

@@ -18,7 +18,7 @@ else:
 
 setup(
     name="conditioned-profile",
-    version="0.1",
+    version="0.2",
     packages=find_packages(),
     description="Firefox Heavy Profile creator",
     include_package_data=True,

testing/condprofile/tox.ini

@@ -5,7 +5,7 @@ envlist = py36,flake8
 [testenv]
 passenv = TRAVIS TRAVIS_JOB_ID TRAVIS_BRANCH
 deps = -rtox-requirements.txt
-       -rrequirements.txt
+       -rlocal-requirements.txt
 commands =
        pytest --random-order-bucket=global -sv --cov-report= --cov-config .coveragerc --cov condprof condprof/tests
        - coverage report -m