Add output_dir arg to download_fuzzer_config script

src/clusterfuzz/_internal/tests/core/local/butler/scripts/download_fuzzer_config_test.py

@@ -15,6 +15,8 @@
 
 import json
 import os
+import shutil
+import tempfile
 import unittest
 
 from clusterfuzz._internal.datastore import data_types
@@ -68,33 +70,61 @@ def setUp(self):
     )
     self.fuzzer2.put()
 
-  def tearDown(self):
-    if os.path.exists('fuzzer1_config.json'):
-      os.remove('fuzzer1_config.json')
-    if os.path.exists('fuzzer2_config.json'):
-      os.remove('fuzzer2_config.json')
+    self.temp_dir = tempfile.mkdtemp()
+    self.addCleanup(shutil.rmtree, self.temp_dir, ignore_errors=True)
+
+  def _config_path(self, fuzzer_name, output_dir=None):
+    return os.path.join(output_dir or self.temp_dir,
+                        f'{fuzzer_name}_config.json')
 
   def test_execute_success(self):
     """Test successful download."""
-    args = Args(['fuzzer1', 'fuzzer2'])
+    args = Args(['fuzzer1', 'fuzzer2', '--output-dir', self.temp_dir])
     download_fuzzer_config.execute(args)
 
-    with open('fuzzer1_config.json') as f:
+    with open(self._config_path('fuzzer1')) as f:
       config1 = json.load(f)
       self.assertEqual(['job1'], config1['jobs'])
       self.assertEqual('data_bundle1', config1['data_bundle_name'])
       self.assertEqual(10, config1['timeout'])
 
-    with open('fuzzer2_config.json') as f:
+    with open(self._config_path('fuzzer2')) as f:
       config2 = json.load(f)
       self.assertEqual([], config2['jobs'])
       self.assertEqual('data_bundle2', config2['data_bundle_name'])
       self.assertEqual(20, config2['timeout'])
 
+  def test_execute_creates_output_dir(self):
+    """Test that a missing nested output directory is created."""
+    output_dir = os.path.join(self.temp_dir, 'fuzzers', 'my_fuzzer',
+                              'clusterfuzz_config')
+    args = Args(['fuzzer1', '--output-dir', output_dir])
+    download_fuzzer_config.execute(args)
+
+    self.assertTrue(os.path.exists(self._config_path('fuzzer1', output_dir)))
+
+  def test_execute_default_output_dir(self):
+    """Test that configs are written to the CWD when --output-dir is unset."""
+    cwd = os.getcwd()
+    os.chdir(self.temp_dir)
+    self.addCleanup(os.chdir, cwd)
+
+    download_fuzzer_config.execute(Args(['fuzzer1']))
+
+    self.assertTrue(os.path.exists(self._config_path('fuzzer1')))
+
+  def test_execute_dry_run(self):
+    """Test that nothing is written in dry-run mode."""
+    output_dir = os.path.join(self.temp_dir, 'dry_run')
+    args = Args(['fuzzer1', '--output-dir', output_dir], non_dry_run=False)
+    download_fuzzer_config.execute(args)
+
+    self.assertFalse(os.path.exists(output_dir))
+
   def test_execute_not_found(self):
     """Test fuzzer not found."""
-    args = Args(['fuzzer1', 'fuzzer_missing'])
+    args = Args(['fuzzer1', 'fuzzer_missing', '--output-dir', self.temp_dir])
     download_fuzzer_config.execute(args)
 
-    self.assertTrue(os.path.exists('fuzzer1_config.json'))
-    self.assertFalse(os.path.exists('fuzzer_missing_config.json'))
+    self.assertTrue(os.path.exists(self._config_path('fuzzer1')))
+    self.assertFalse(os.path.exists(self._config_path('fuzzer_missing')))

src/local/butler/scripts/download_fuzzer_config.py

@@ -13,34 +13,56 @@
 # limitations under the License.
 """Download fuzzer config as a JSON file."""
 
+import argparse
 import json
+import os
 import sys
 
 from clusterfuzz._internal.datastore import data_types
 
 
+def _parse_script_args(script_args):
+  """Parses the script specific arguments."""
+  parser = argparse.ArgumentParser(prog='download_fuzzer_config')
+  parser.add_argument(
+      'fuzzer_names', nargs='+', help='Names of the fuzzers to download.')
+  parser.add_argument(
+      '--output-dir',
+      default='.',
+      help='Directory to write <fuzzer_name>_config.json files to. Created if '
+      'it does not exist. Defaults to the current working directory.')
+  return parser.parse_args(script_args)
+
+
 def execute(args):
   """Download fuzzer config."""
   if not args.script_args:
     print('Please provide a list of fuzzer names as script arguments.')
     sys.exit(1)
 
+  script_args = _parse_script_args(args.script_args)
+  fuzzer_names = script_args.fuzzer_names
+  output_dir = script_args.output_dir
+
   fuzzers = data_types.Fuzzer.query(
-      data_types.Fuzzer.name.IN(args.script_args)).fetch()
+      data_types.Fuzzer.name.IN(fuzzer_names)).fetch()
 
   existing_fuzzer_names = {fuzzer.name for fuzzer in fuzzers}
 
-  for fuzzer_name in args.script_args:
+  for fuzzer_name in fuzzer_names:
     if fuzzer_name not in existing_fuzzer_names:
       print(f'Fuzzer {fuzzer_name} not found.')
 
   if not args.non_dry_run:
     print('Skipping writes in dry-run mode.')
     return
 
+  if fuzzers:
+    os.makedirs(output_dir, exist_ok=True)
+
   for fuzzer in fuzzers:
     config = fuzzer.get_config_dict()
-    filename = f'{fuzzer.name}_config.json'
+    filename = os.path.join(output_dir, f'{fuzzer.name}_config.json')
 
     with open(filename, 'w') as f:
       json.dump(config, f, indent=4)