diff --git a/testing/condprofile/condprof/client.py b/testing/condprofile/condprof/client.py index 1d71acf2da36..cea8ea9bd5bb 100644 --- a/testing/condprofile/condprof/client.py +++ b/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 - download_dir = tempfile.mkdtemp() + 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) - archive = download_file(url, target=downloaded_archive) try: with tarfile.open(archive, "r:gz") as tar: get_logger().msg("Extracting the tarball content in %s" % target_dir) @@ -76,7 +91,8 @@ def get_profile(target_dir, platform, scenario, customization="default", task_id except (OSError, tarfile.ReadError) as e: raise ProfileNotFoundError(str(e)) finally: - shutil.rmtree(download_dir) + 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") - download_file(changelog_url, target=downloaded_changelog) + try: + download_file(changelog_url, target=downloaded_changelog) + except ArchiveNotFound: + shutil.rmtree(download_dir) + raise ProfileNotFoundError(changelog_url) return Changelog(download_dir) diff --git a/testing/condprofile/condprof/metadata.py b/testing/condprofile/condprof/metadata.py index 54ef806cfe54..ab338dd72f13 100644 --- a/testing/condprofile/condprof/metadata.py +++ b/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. """ diff --git a/testing/condprofile/condprof/tests/test_client.py b/testing/condprofile/condprof/tests/test_client.py new file mode 100644 index 000000000000..185156262ee5 --- /dev/null +++ b/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) diff --git a/testing/condprofile/condprof/tests/test_runner.py b/testing/condprofile/condprof/tests/test_runner.py index 4c3a69bf50a3..4295ffe3b687 100644 --- a/testing/condprofile/condprof/tests/test_runner.py +++ b/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( diff --git a/testing/condprofile/setup.py b/testing/condprofile/setup.py index 85914058224b..c7122dd78139 100644 --- a/testing/condprofile/setup.py +++ b/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, diff --git a/testing/condprofile/tox.ini b/testing/condprofile/tox.ini index 316b804a7870..40cebf9fdf3f 100644 --- a/testing/condprofile/tox.ini +++ b/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