src/third_party/chromite/lib/paygen/download_cache_unittest.py

   1 #!/usr/bin/python
   2 # Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
   3 # Use of this source code is governed by a BSD-style license that can be
   4 # found in the LICENSE file.
   5
   6 """Test download_cache library.
   7
   8 DEPRECATED: Should be migrated to chromite.lib.cache_unittest.
   9 """
  10
  11 from __future__ import print_function
  12
  13 import mox
  14 import multiprocessing
  15 import os
  16 import pickle
  17 import traceback
  18
  19 import fixup_path
  20 fixup_path.FixupPath()
  21
  22 from chromite.lib import cros_test_lib
  23 from chromite.lib import osutils
  24 from chromite.lib.paygen import download_cache
  25 from chromite.lib.paygen import gslib
  26 from chromite.lib.paygen import unittest_lib
  27
  28
  29 # We access a lot of protected members during testing.
  30 # pylint: disable-msg=W0212
  31
  32 # The inProcess methods have to be standalone to be pickleable.
  33 def _inProcessFetchIntoCache(uri_tempdir):
  34   """In a sub-process, call DownloadCache._UriToCacheFile."""
  35   try:
  36     uri, tempdir = uri_tempdir
  37     process_cache = download_cache.DownloadCache(tempdir)
  38     file_name = process_cache._UriToCacheFile(uri)
  39     with process_cache._PurgeLock(shared=True, blocking=True):
  40       return process_cache._FetchIntoCache(uri, file_name)
  41   except Exception:
  42     traceback.print_exc()
  43     raise
  44
  45
  46 def _inProcessGetFile(uri_tempdir):
  47   """In a sub-process, call DownloadCache.GetFile."""
  48
  49   try:
  50     uri, tempdir = uri_tempdir
  51     process_cache = download_cache.DownloadCache(tempdir, cache_size=0)
  52
  53     # If there is a URI, fetch it, else wipe.
  54     if uri:
  55       with process_cache.GetFileObject(uri) as f:
  56         return f.read()
  57     else:
  58       process_cache.Purge()
  59       return None
  60   except Exception:
  61     traceback.print_exc()
  62     raise
  63
  64
  65 class DownloadCachePickleTest(unittest_lib.TestCase):
  66   """Test pickle/unpickle the download cache."""
  67
  68   @osutils.TempDirDecorator
  69   def testPickleUnpickle(self):
  70     # pylint: disable-msg=E1101
  71     cache = download_cache.DownloadCache(self.tempdir)
  72     pickle_path = os.path.join(self.tempdir, 'cache.pickle')
  73
  74     # Do pickle dump.
  75     with open(pickle_path, 'w') as pickle_fh:
  76       pickle.dump(cache, pickle_fh)
  77
  78     # Load pickle file.
  79     with open(pickle_path, 'r') as pickle_fh:
  80       pickle.load(pickle_fh)
  81
  82
  83 class DownloadCacheTest(mox.MoxTestBase):
  84   """Test DownloadCache helper class."""
  85
  86   def __init__(self, testCaseNames):
  87     self.tempdir = None
  88     mox.MoxTestBase.__init__(self, testCaseNames)
  89
  90     self.uri_large = 'gs://chromeos-releases-test/download_cache/file_large'
  91     self.uri_a = 'gs://chromeos-releases-test/download_cache/file_a'
  92     self.uri_b = 'gs://chromeos-releases-test/download_cache/file_b'
  93
  94
  95     self.hash_large = 'ce11166b2742c12c93efa307c4c4adbf'
  96     self.hash_a = '591430f83b55355d9233babd172baea5'
  97     self.hash_b = '22317eb6cccea8c87f960c45ecec3478'
  98
  99   def setUp(self):
 100     """Prepare for each test."""
 101     self.mox = mox.Mox()
 102
 103     # To make certain we don't self update while running tests.
 104     os.environ['CROSTOOLS_NO_SOURCE_UPDATE'] = '1'
 105
 106   def tearDown(self):
 107     """Cleanup after each test."""
 108     self.mox.UnsetStubs()
 109
 110   def _verifyFileContents(self, cache, uri):
 111     """Test helper to make sure a cached file contains correct contents."""
 112
 113     # Fetch it
 114     with cache.GetFileObject(uri) as f:
 115       contents = f.read()
 116
 117     # Make sure the contents are valid.
 118     self.assertEqual(contents, gslib.Cat(uri))
 119
 120     # Make sure the cache file exists where expected.
 121     cache_file = cache._UriToCacheFile(uri)
 122
 123     self.assertTrue(cache_file.startswith(self.tempdir))
 124     self.assertTrue(os.path.exists(cache_file))
 125
 126   def _validateCacheContents(self, cache, expected_contents):
 127     """Test helper to make sure the cache holds what we expect."""
 128
 129     expected_contents = set(expected_contents)
 130     expected_top_contents = set(['cache', 'cache.lock', 'lock'])
 131
 132     cache_top_contents = set(os.listdir(cache._cache_dir))
 133     file_dir_contents = set(os.listdir(cache._file_dir))
 134     lock_dir_contents = set(os.listdir(cache._lock_dir))
 135
 136     # We should always have exactly the expected files in the top dir.
 137     self.assertEqual(cache_top_contents, expected_top_contents)
 138
 139     # Cache contents should match the expected list.
 140     self.assertEqual(file_dir_contents, expected_contents)
 141
 142     # The lock directory should contain no files not in the file_dir.
 143     self.assertTrue(lock_dir_contents.issubset(file_dir_contents))
 144
 145   @osutils.TempDirDecorator
 146   def testCacheFileNames(self):
 147     """Make sure that some of the files we create have the expected names."""
 148     cache = download_cache.DownloadCache(self.tempdir)
 149
 150     expected_cache_lock = os.path.join(self.tempdir, 'cache.lock')
 151     expected_cache = os.path.join(self.tempdir,
 152                                   'cache/3ba505fc7774455169af6f50b7964dff')
 153
 154     expected_lock = os.path.join(self.tempdir,
 155                                  'lock/3ba505fc7774455169af6f50b7964dff')
 156
 157     # Make sure a cache content file is named as expected.
 158     self.assertEqual(cache._UriToCacheFile('gs://bucket/of/awesome'),
 159                      expected_cache)
 160
 161     # Make sure the lock file for a cache content file is named as expected.
 162     file_lock = cache._CacheFileLock(expected_cache)
 163     self.assertEqual(file_lock.lock_file, expected_lock)
 164
 165     purge_lock = cache._PurgeLock()
 166     self.assertEqual(purge_lock.lock_file, expected_cache_lock)
 167
 168     cache_file_lock = cache._CacheFileLock(expected_cache)
 169     self.assertEqual(cache_file_lock.lock_file, expected_lock)
 170
 171   @osutils.TempDirDecorator
 172   def testSetupCacheClean(self):
 173     """Test _SetupCache with a clean directory."""
 174     # Create a cache, and see if it has expected contents.
 175     cache = download_cache.DownloadCache(self.tempdir)
 176     self._validateCacheContents(cache, ())
 177
 178   @osutils.TempDirDecorator
 179   def testSetupCacheDirty(self):
 180     """Test _SetupCache with a dirty directory."""
 181     # Create some unexpected directories.
 182     for make_dir in ['foo/bar/stuff', 'bar']:
 183       os.makedirs(os.path.join(self.tempdir, make_dir))
 184
 185     # Touch some unexpected files.
 186     for touch_file in ['bogus', 'foo/bogus']:
 187       file(os.path.join(self.tempdir, touch_file), 'w').close()
 188
 189     # Create a cache, and see
 190     cache = download_cache.DownloadCache(self.tempdir)
 191     self._validateCacheContents(cache, ())
 192
 193   @cros_test_lib.NetworkTest()
 194   @osutils.TempDirDecorator
 195   def testGetFileObject(self):
 196     """Just create a download cache, and GetFile on it."""
 197
 198     cache = download_cache.DownloadCache(self.tempdir)
 199
 200     # Fetch a file
 201     with cache.GetFileObject(self.uri_a) as f:
 202       self.assertIsInstance(f, file)
 203     self._verifyFileContents(cache, self.uri_a)
 204     self._validateCacheContents(cache, (self.hash_a,))
 205
 206     # Fetch a different file
 207     with cache.GetFileObject(self.uri_b) as f:
 208       self.assertIsInstance(f, file)
 209     self._verifyFileContents(cache, self.uri_b)
 210     self._validateCacheContents(cache, (self.hash_a, self.hash_b))
 211
 212     # Fetch the first file a second time.
 213     cache.GetFileObject(self.uri_a).close()
 214     self._verifyFileContents(cache, self.uri_a)
 215
 216     # There should be only 2 files in the cache.
 217     self._validateCacheContents(cache, (self.hash_a, self.hash_b))
 218
 219     # Fetch a larger file
 220     cache.GetFileObject(self.uri_large).close()
 221     self._verifyFileContents(cache, self.uri_large)
 222
 223     # There should be 3 files in the cache.
 224     self._validateCacheContents(cache,
 225                                 (self.hash_a, self.hash_b, self.hash_large))
 226
 227   @cros_test_lib.NetworkTest()
 228   @osutils.TempDirDecorator
 229   def testGetFileCopy(self):
 230     """Just create a download cache, and GetFileCopy from it."""
 231
 232     file_a = os.path.join(self.tempdir, 'foo')
 233     file_b = os.path.join(self.tempdir, 'bar')
 234     cache_dir = os.path.join(self.tempdir, 'cache')
 235
 236     cache = download_cache.DownloadCache(cache_dir)
 237
 238     # Fetch non-existent files.
 239     cache.GetFileCopy(self.uri_a, file_a)
 240     cache.GetFileCopy(self.uri_a, file_b)
 241
 242     with open(file_a, 'r') as f:
 243       contents_a = f.read()
 244
 245     with open(file_b, 'r') as f:
 246       contents_b = f.read()
 247
 248     self.assertEqual(contents_a, contents_b)
 249
 250     # Fetch and overwrite existent files.
 251     cache.GetFileCopy(self.uri_b, file_a)
 252     cache.GetFileCopy(self.uri_b, file_b)
 253
 254     with open(file_a, 'r') as f:
 255       contents_a = f.read()
 256
 257     with open(file_b, 'r') as f:
 258       contents_b = f.read()
 259
 260     self.assertEqual(contents_a, contents_b)
 261
 262   @cros_test_lib.NetworkTest()
 263   @osutils.TempDirDecorator
 264   def testGetFileInTempFile(self):
 265     """Just create a download cache, and GetFileInTempFile on it."""
 266
 267     cache = download_cache.DownloadCache(self.tempdir)
 268
 269     # Fetch a file
 270     file_t = cache.GetFileInTempFile(self.uri_a)
 271
 272     with cache.GetFileObject(self.uri_a) as f:
 273       contents_a = f.read()
 274
 275     with file_t as f:
 276       contents_t = f.read()
 277
 278     self.assertEqual(contents_t, contents_a)
 279     self.assertEqual(contents_t, gslib.Cat(self.uri_a))
 280
 281   @cros_test_lib.NetworkTest()
 282   @osutils.TempDirDecorator
 283   def testPurgeLogic(self):
 284     cache = download_cache.DownloadCache(self.tempdir)
 285
 286     cache.GetFileObject(self.uri_a).close()
 287     cache.GetFileObject(self.uri_b).close()
 288
 289     # The default cache logic should leave these files untouched, since
 290     # they are less than a day old.
 291     cache.Purge()
 292     self._validateCacheContents(cache, (self.hash_a, self.hash_b))
 293
 294     # Purge until the cache is empty.
 295     cache.Purge(cache_size=0)
 296     self._validateCacheContents(cache, ())
 297
 298     # Refetch two files.
 299     cache.GetFileObject(self.uri_a).close()
 300     cache.GetFileObject(self.uri_b).close()
 301
 302     # Change the timestamp so uri_a hasn't been used for a very long time.
 303     os.utime(os.path.join(self.tempdir, 'cache', self.hash_a),
 304              (2, 2))
 305
 306     # Purge files that haven't been used recently.
 307     cache.Purge(max_age=1000)
 308     self._validateCacheContents(cache, (self.hash_b,))
 309
 310   @cros_test_lib.NetworkTest()
 311   @osutils.TempDirDecorator
 312   def testContextMgr(self):
 313     """Make sure we behave properly with 'with'."""
 314
 315     # Create an instance, and use it in a with
 316     precache = download_cache.DownloadCache(self.tempdir, cache_size=0)
 317
 318     with precache as cache:
 319       # Assert the instance didn't change.
 320       self.assertIs(precache, cache)
 321
 322       # Download a file.
 323       cache.GetFileObject(self.uri_a).close()
 324
 325       self._validateCacheContents(cache, (self.hash_a,))
 326
 327     # After the with exited, which should have purged everything.
 328     self._validateCacheContents(cache, ())
 329
 330   @cros_test_lib.NetworkTest()
 331   @osutils.TempDirDecorator
 332   def testThreadedDownloads(self):
 333     """Spin off multiple processes and fetch a file.
 334
 335        Ensure the process locking allows the file to be downloaded exactly
 336        once.
 337     """
 338     pool = multiprocessing.Pool(processes=10)
 339
 340     # Create a tuple of the three args we want to pass to inProcess test,
 341     # use map semantics as a convenient way to run in parallel.
 342     results = pool.map(_inProcessFetchIntoCache,
 343                        [(self.uri_large, self.tempdir)] * 20)
 344
 345     # Results contains a list of booleans showing which instances actually
 346     # performed the download. Exactly one of them should have. The list could
 347     # also contain exceptions if one of the downloads failed.
 348     results.sort()
 349     self.assertEqual(results, [False] * 19 + [True])
 350
 351   @cros_test_lib.NetworkTest()
 352   @osutils.TempDirDecorator
 353   def testThreadedGetFile(self):
 354     """Spin off multiple processes and call GetFile.
 355
 356        Ensure all processes complete, and return the same local file.
 357     """
 358     pool = multiprocessing.Pool(processes=10)
 359
 360     # Create a tuple of the three args we want to pass to inProcess test,
 361     # use map semantics as a convenient way to run in parallel.
 362     results = pool.map(_inProcessGetFile,
 363                        [(self.uri_a, self.tempdir)] * 20)
 364
 365     # Fetch it ourselves and verify the results.
 366     cache = download_cache.DownloadCache(self.tempdir)
 367     self._verifyFileContents(cache, self.uri_a)
 368
 369     with cache.GetFileObject(self.uri_a) as f:
 370       contents_a = f.read()
 371
 372     # Ensure that every process gave back the expected result.
 373     expected = [contents_a] * 20
 374     self.assertEqual(results, expected)
 375
 376   @cros_test_lib.NetworkTest()
 377   @osutils.TempDirDecorator
 378   def testThreadedGetFileMultiple(self):
 379     """Spin off multiple processes and call GetFile with multiple uris.
 380
 381        Ensure all processes complete, and return the right local file.
 382     """
 383     pool = multiprocessing.Pool(processes=20)
 384
 385     # Create a tuple of the three args we want to pass to inProcess test,
 386     # use map semantics as a convenient way to run in parallel.
 387     results = pool.map(_inProcessGetFile,
 388                        [(self.uri_a, self.tempdir),
 389                         (self.uri_b, self.tempdir)] * 10)
 390
 391     # Fetch it ourselves and verify the results.
 392     cache = download_cache.DownloadCache(self.tempdir)
 393
 394     with cache.GetFileObject(self.uri_a) as f:
 395       contents_a = f.read()
 396
 397     with cache.GetFileObject(self.uri_b) as f:
 398       contents_b = f.read()
 399
 400     self._verifyFileContents(cache, self.uri_a)
 401     self._verifyFileContents(cache, self.uri_b)
 402
 403     # Ensure that every process gave back the expected result.
 404     expected = [contents_a, contents_b] * 10
 405     self.assertEqual(results, expected)
 406
 407   @cros_test_lib.NetworkTest()
 408   @osutils.TempDirDecorator
 409   def testThreadedGetFileMultiplePurge(self):
 410     """Do fetches and purges in a multiprocess environment.
 411
 412        Ensure all processes complete, and return the right local file.
 413     """
 414     pool = multiprocessing.Pool(processes=30)
 415
 416     requests = [(self.uri_a, self.tempdir),
 417                 (self.uri_b, self.tempdir),
 418                 (None, self.tempdir)] * 10
 419
 420     # Create a tuple of the three args we want to pass to inProcess test,
 421     # use map semantics as a convenient way to run in parallel.
 422     results = pool.map(_inProcessGetFile, requests)
 423
 424     # Fetch it ourselves and verify the results.
 425     cache = download_cache.DownloadCache(self.tempdir)
 426
 427     with cache.GetFileObject(self.uri_a) as f:
 428       contents_a = f.read()
 429
 430     with cache.GetFileObject(self.uri_b) as f:
 431       contents_b = f.read()
 432
 433     self._verifyFileContents(cache, self.uri_a)
 434     self._verifyFileContents(cache, self.uri_b)
 435
 436     # Ensure that every process gave back the expected result.
 437     expected = [contents_a, contents_b, None] * 10
 438     self.assertEqual(results, expected)
 439
 440
 441 if __name__ == '__main__':
 442   cros_test_lib.main()