python/mozbuild/mozbuild/test/test_artifact_cache.py

   1 # This Source Code Form is subject to the terms of the Mozilla Public
   2 # License, v. 2.0. If a copy of the MPL was not distributed with this
   3 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
   4
   5 import os
   6 import time
   7 import unittest
   8 from shutil import rmtree
   9 from tempfile import mkdtemp
  10
  11 import mozunit
  12
  13 from mozbuild import artifact_cache
  14 from mozbuild.artifact_cache import ArtifactCache
  15
  16 CONTENTS = {
  17     "http://server/foo": b"foo",
  18     "http://server/bar": b"bar" * 400,
  19     "http://server/qux": b"qux" * 400,
  20     "http://server/fuga": b"fuga" * 300,
  21     "http://server/hoge": b"hoge" * 300,
  22     "http://server/larger": b"larger" * 3000,
  23 }
  24
  25
  26 class FakeResponse(object):
  27     def __init__(self, content):
  28         self._content = content
  29
  30     @property
  31     def headers(self):
  32         return {"Content-length": str(len(self._content))}
  33
  34     def iter_content(self, chunk_size):
  35         content = memoryview(self._content)
  36         while content:
  37             yield content[:chunk_size]
  38             content = content[chunk_size:]
  39
  40     def raise_for_status(self):
  41         pass
  42
  43     def close(self):
  44         pass
  45
  46
  47 class FakeSession(object):
  48     def get(self, url, stream=True):
  49         assert stream is True
  50         return FakeResponse(CONTENTS[url])
  51
  52
  53 class TestArtifactCache(unittest.TestCase):
  54     def setUp(self):
  55         self.min_cached_artifacts = artifact_cache.MIN_CACHED_ARTIFACTS
  56         self.max_cached_artifacts_size = artifact_cache.MAX_CACHED_ARTIFACTS_SIZE
  57         artifact_cache.MIN_CACHED_ARTIFACTS = 2
  58         artifact_cache.MAX_CACHED_ARTIFACTS_SIZE = 4096
  59
  60         self._real_utime = os.utime
  61         os.utime = self.utime
  62         self.timestamp = time.time() - 86400
  63
  64         self.tmpdir = mkdtemp()
  65
  66     def tearDown(self):
  67         rmtree(self.tmpdir)
  68         artifact_cache.MIN_CACHED_ARTIFACTS = self.min_cached_artifacts
  69         artifact_cache.MAX_CACHED_ARTIFACTS_SIZE = self.max_cached_artifacts_size
  70         os.utime = self._real_utime
  71
  72     def utime(self, path, times):
  73         if times is None:
  74             # Ensure all downloaded files have a different timestamp
  75             times = (self.timestamp, self.timestamp)
  76             self.timestamp += 2
  77         self._real_utime(path, times)
  78
  79     def listtmpdir(self):
  80         return [p for p in os.listdir(self.tmpdir) if p != ".metadata_never_index"]
  81
  82     def test_artifact_cache_persistence(self):
  83         cache = ArtifactCache(self.tmpdir)
  84         cache._download_manager.session = FakeSession()
  85
  86         path = cache.fetch("http://server/foo")
  87         expected = [os.path.basename(path)]
  88         self.assertEqual(self.listtmpdir(), expected)
  89
  90         path = cache.fetch("http://server/bar")
  91         expected.append(os.path.basename(path))
  92         self.assertEqual(sorted(self.listtmpdir()), sorted(expected))
  93
  94         # We're downloading more than the cache allows us, but since it's all
  95         # in the same session, no purge happens.
  96         path = cache.fetch("http://server/qux")
  97         expected.append(os.path.basename(path))
  98         self.assertEqual(sorted(self.listtmpdir()), sorted(expected))
  99
 100         path = cache.fetch("http://server/fuga")
 101         expected.append(os.path.basename(path))
 102         self.assertEqual(sorted(self.listtmpdir()), sorted(expected))
 103
 104         cache = ArtifactCache(self.tmpdir)
 105         cache._download_manager.session = FakeSession()
 106
 107         # Downloading a new file in a new session purges the oldest files in
 108         # the cache.
 109         path = cache.fetch("http://server/hoge")
 110         expected.append(os.path.basename(path))
 111         expected = expected[2:]
 112         self.assertEqual(sorted(self.listtmpdir()), sorted(expected))
 113
 114         # Downloading a file already in the cache leaves the cache untouched
 115         cache = ArtifactCache(self.tmpdir)
 116         cache._download_manager.session = FakeSession()
 117
 118         path = cache.fetch("http://server/qux")
 119         self.assertEqual(sorted(self.listtmpdir()), sorted(expected))
 120
 121         # bar was purged earlier, re-downloading it should purge the oldest
 122         # downloaded file, which at this point would be qux, but we also
 123         # re-downloaded it in the mean time, so the next one (fuga) should be
 124         # the purged one.
 125         cache = ArtifactCache(self.tmpdir)
 126         cache._download_manager.session = FakeSession()
 127
 128         path = cache.fetch("http://server/bar")
 129         expected.append(os.path.basename(path))
 130         expected = [p for p in expected if "fuga" not in p]
 131         self.assertEqual(sorted(self.listtmpdir()), sorted(expected))
 132
 133         # Downloading one file larger than the cache size should still leave
 134         # MIN_CACHED_ARTIFACTS files.
 135         cache = ArtifactCache(self.tmpdir)
 136         cache._download_manager.session = FakeSession()
 137
 138         path = cache.fetch("http://server/larger")
 139         expected.append(os.path.basename(path))
 140         expected = expected[-2:]
 141         self.assertEqual(sorted(self.listtmpdir()), sorted(expected))
 142
 143
 144 if __name__ == "__main__":
 145     mozunit.main()