Add CIPD support for bot assets

infra/bots/assets/asset_utils_test.py

Index: infra/bots/assets/asset_utils_test.py
diff --git a/infra/bots/assets/asset_utils_test.py b/infra/bots/assets/asset_utils_test.py
index edfc271f9e769e378d1173865f0daa7ee870842e..a76042fe770be1e8e905764b5ab7315c62c9674c 100644
--- a/infra/bots/assets/asset_utils_test.py
+++ b/infra/bots/assets/asset_utils_test.py
@@ -20,13 +20,13 @@ import uuid
 
 
 FILE_DIR = os.path.dirname(os.path.abspath(__file__))
-INFRA_BOTS_DIR = os.path.realpath(os.path.join(
-    FILE_DIR, os.pardir, 'infra', 'bots'))
+INFRA_BOTS_DIR = os.path.realpath(os.path.join(FILE_DIR, os.pardir))
 sys.path.insert(0, INFRA_BOTS_DIR)
 import test_utils
 import utils
 
 
+CIPD_DEV_SERVICE_URL = 'https://chrome-infra-packages-dev.appspot.com'
 GS_BUCKET = 'skia-infra-testdata'
 
 
@@ -48,16 +48,131 @@ def _write_stuff(target_dir):
   fw.write(os.path.join('subdir', 'd.txt'), 0640)
 
 
-class AssetUtilsTest(unittest.TestCase):
+class _LocalStore(object):
+  """Local store used for testing."""
+  def __init__(self):
+    self.dir = tempfile.mkdtemp()
+
+  def get_available_versions(self, name):
+    target = os.path.join(self.dir, name)
+    if not os.path.isdir(target):
+      return []
+    contents = os.listdir(os.path.join(self.dir, name))
+    return sorted([int(d) for d in contents])
+
+  def upload(self, name, version, target_dir):
+    shutil.copytree(target_dir, os.path.join(self.dir, name, str(version)))
+
+  def download(self, name, version, target_dir):
+    shutil.copytree(os.path.join(self.dir, name, str(version)), target_dir)
+
+  def delete_contents(self, name):
+    try:
+      shutil.rmtree(self.dir)
+    except OSError:
+      if os.path.exists(self.dir):
+        raise
+
+
+class StoreTest(unittest.TestCase):
+  """Superclass used for testing one of the stores."""
+  def setUp(self):
+    self.asset_name = str(uuid.uuid4())
+
+  def tearDown(self):
+    pass
+
+  def _test_upload_download(self, store):
+    with utils.tmp_dir():
+      # Create input files and directories.
+      input_dir = os.path.join(os.getcwd(), 'input')
+      _write_stuff(input_dir)
+
+      # Upload a version, download it again.
+      store.upload(self.asset_name, 0, input_dir)
+      output_dir = os.path.join(os.getcwd(), 'output')
+      store.download(self.asset_name, 0, output_dir)
+
+      # Compare.
+      test_utils.compare_trees(self, input_dir, output_dir)
+
+  def _test_versions(self, store):
+    with utils.tmp_dir():
+      # Create input files and directories.
+      input_dir = os.path.join(os.getcwd(), 'input')
+      _write_stuff(input_dir)
+      self.assertEqual(store.get_available_versions(self.asset_name), [])
+      store.upload(self.asset_name, 0, input_dir)
+      self.assertEqual(store.get_available_versions(self.asset_name), [0])
+      store.upload(self.asset_name, 1, input_dir)
+      self.assertEqual(store.get_available_versions(self.asset_name), [0, 1])
+      store.delete_contents(self.asset_name)
+      self.assertEqual(store.get_available_versions(self.asset_name), [])
+
+
+class LocalStoreTest(StoreTest):
+  """Test the local store."""
+  def setUp(self):
+    super(LocalStoreTest, self).setUp()
+    self._store = _LocalStore()
+
+  def tearDown(self):
+    self._store.delete_contents(self.asset_name)
+    super(LocalStoreTest, self).tearDown()
+
+  def test_upload_download(self):
+    self._test_upload_download(self._store)
+
+  def test_versions(self):
+    self._test_versions(self._store)
+
+
+class CIPDStoreTest(StoreTest):
+  """Test the CIPD store."""
+  def setUp(self):
+    super(CIPDStoreTest, self).setUp()
+    self._store = asset_utils.CIPDStore(cipd_url=CIPD_DEV_SERVICE_URL)
+
+  def tearDown(self):
+    self._store.delete_contents(self.asset_name)
+    super(CIPDStoreTest, self).tearDown()
+
+  def test_upload_download(self):
+    self._test_upload_download(self._store)
+
+  def test_versions(self):
+    self._test_versions(self._store)
+
+
+class GSStoreTest(StoreTest):
+  """Test the GS store."""
+  def setUp(self):
+    super(GSStoreTest, self).setUp()
+    self._store = asset_utils.GSStore(gsutil=None, bucket=GS_BUCKET)
+
+  def tearDown(self):
+    self._store.delete_contents(self.asset_name)
+    super(GSStoreTest, self).tearDown()
+
+  def test_upload_download(self):
+    self._test_upload_download(self._store)
+
+  def test_versions(self):
+    self._test_versions(self._store)
+
+
+class AssetTest(unittest.TestCase):
+  """Test Asset operations using a local store."""
   def setUp(self):
     self.asset_name = str(uuid.uuid4())
     self.old_prompt = asset_utils._prompt
     asset_utils._prompt = _fake_prompt('y')
-    self.a = asset_utils.Asset.add(self.asset_name, gs_bucket=GS_BUCKET)
+    self._store = _LocalStore()
+    self.a = asset_utils.Asset.add(self.asset_name, self._store)
 
   def tearDown(self):
     if self.a:
-      self.a.remove()
+      self.a.remove(remove_in_store=True)
     asset_utils._prompt = self.old_prompt
 
     gs_path = 'gs://%s/assets/%s' % (GS_BUCKET, self.asset_name)
@@ -72,7 +187,7 @@ class AssetUtilsTest(unittest.TestCase):
   def test_add_remove(self):
     # Ensure that we can't create an asset twice.
     with self.assertRaises(Exception):
-      asset_utils.Asset.add(self.asset_name, gs_bucket=GS_BUCKET)
+      asset_utils.Asset.add(self.asset_name, self._store)
 
     # Ensure that the asset dir exists.
     asset_dir = os.path.join(FILE_DIR, self.asset_name)