1 files changed, 32 insertions, 42 deletions
diff --git a/infra/cifuzz/clusterfuzz_deployment_test.py b/infra/cifuzz/clusterfuzz_deployment_test.py
index a7c088126..06ff78476 100644
--- a/infra/cifuzz/clusterfuzz_deployment_test.py
+++ b/infra/cifuzz/clusterfuzz_deployment_test.py
@@ -14,7 +14,6 @@
 """Tests for clusterfuzz_deployment.py"""
 
 import os
-import tempfile
 import unittest
 from unittest import mock
 import urllib.error
@@ -58,47 +57,39 @@ def _create_deployment(**kwargs):
   return clusterfuzz_deployment.get_clusterfuzz_deployment(config)
 
 
-class OSSFuzzTest(unittest.TestCase):
+class OSSFuzzTest(fake_filesystem_unittest.TestCase):
   """Tests OSSFuzz."""
 
-  def test_download_corpus(self):
+  OUT_DIR = '/out'
+
+  def setUp(self):
+    self.setUpPyfakefs()
+    self.deployment = _create_deployment()
+
+  @mock.patch('clusterfuzz_deployment.download_and_unpack_zip',
+              return_value=True)
+  def test_download_corpus(self, mocked_download_and_unpack_zip):
     """Tests that we can download a corpus for a valid project."""
-    deployment = _create_deployment()
-    with tempfile.TemporaryDirectory() as tmp_dir:
-      with mock.patch('clusterfuzz_deployment.download_and_unpack_zip',
-                      return_value=False) as mocked_download_and_unpack_zip:
-        deployment.download_corpus(EXAMPLE_FUZZER, tmp_dir)
-        (url, out_dir), _ = mocked_download_and_unpack_zip.call_args
-        self.assertEqual(
-            url, 'https://storage.googleapis.com/example-backup.'
-            'clusterfuzz-external.appspot.com/corpus/libFuzzer/'
-            'example_crash_fuzzer/public.zip')
-        self.assertEqual(out_dir,
-                         os.path.join(tmp_dir, 'cifuzz-corpus', EXAMPLE_FUZZER))
-
-  def test_download_fail(self):
+    result = self.deployment.download_corpus(EXAMPLE_FUZZER, self.OUT_DIR)
+    self.assertIsNotNone(result)
+    expected_corpus_dir = os.path.join(self.OUT_DIR, 'cifuzz-corpus',
+                                       EXAMPLE_FUZZER)
+    expected_url = ('https://storage.googleapis.com/example-backup.'
+                    'clusterfuzz-external.appspot.com/corpus/libFuzzer/'
+                    'example_crash_fuzzer/public.zip')
+    call_args, _ = mocked_download_and_unpack_zip.call_args
+    self.assertEqual(call_args, (expected_url, expected_corpus_dir))
+
+  @mock.patch('clusterfuzz_deployment.download_and_unpack_zip',
+              return_value=False)
+  def test_download_fail(self, _):
     """Tests that when downloading fails, None is returned."""
-    deployment = _create_deployment()
-    with tempfile.TemporaryDirectory() as tmp_dir:
-      with mock.patch('clusterfuzz_deployment.download_and_unpack_zip',
-                      return_value=False):
-        corpus_path = deployment.download_corpus(EXAMPLE_FUZZER, tmp_dir)
-        self.assertIsNone(corpus_path)
-
-  def test_download_latest_build(self):
-    """Tests that the build directory is downloaded once and no more."""
-    deployment = _create_deployment()
-    with tempfile.TemporaryDirectory() as tmp_dir:
-      latest_name = deployment.get_latest_build_name()
-      with mock.patch('clusterfuzz_deployment.OSSFuzz.get_latest_build_name',
-                      return_value=latest_name):
-        latest_build_path = deployment.download_latest_build(tmp_dir)
-        self.assertNotEqual(len(os.listdir(latest_build_path)), 0)
+    corpus_path = self.deployment.download_corpus(EXAMPLE_FUZZER, self.OUT_DIR)
+    self.assertIsNone(corpus_path)
 
   def test_get_latest_build_name(self):
     """Tests that the latest build name can be retrieved from GCS."""
-    deployment = _create_deployment()
-    latest_build_name = deployment.get_latest_build_name()
+    latest_build_name = self.deployment.get_latest_build_name()
     self.assertTrue(latest_build_name.endswith('.zip'))
     self.assertTrue('address' in latest_build_name)
 
@@ -147,14 +138,13 @@ class DownloadAndUnpackZipTest(fake_filesystem_unittest.TestCase):
   def setUp(self):
     self.setUpPyfakefs()
 
-  def test_bad_zip_download(self):
+  @mock.patch('urllib.request.urlretrieve', return_value=True)
+  def test_bad_zip_download(self, _):
     """Tests download_and_unpack_zip returns none when a bad zip is passed."""
-    with open('/url_tmp.zip', 'w') as file_handle:
-      file_handle.write('Test file.')
-    with mock.patch('urllib.request.urlretrieve', return_value=True):
-      self.assertFalse(
-          clusterfuzz_deployment.download_and_unpack_zip(
-              '/not/a/real/url', '/extract-directory'))
+    self.fs.create_file('/url_tmp.zip', contents='Test file.')
+    self.assertFalse(
+        clusterfuzz_deployment.download_and_unpack_zip('/not/a/real/url',
+                                                       '/extract-directory'))
 
 
 if __name__ == '__main__':