tf_fuzz: suitegen add option to generate multiple tests per template

Add -n option to tfz-suitegen, allowing multiple test cases to be
generated for each test template file.

Change-Id: If74f65d04ba9b5ef1ce27b7c51c4e1a02e62b51b
Signed-off-by: Nik Dewally <Nik.Dewally@arm.com>
diff --git a/tf_fuzz/tfz-suitegen/src/tfz-suitegen/__main__.py b/tf_fuzz/tfz-suitegen/src/tfz-suitegen/__main__.py
index 49340f4..be09f36 100644
--- a/tf_fuzz/tfz-suitegen/src/tfz-suitegen/__main__.py
+++ b/tf_fuzz/tfz-suitegen/src/tfz-suitegen/__main__.py
@@ -49,8 +49,8 @@
 
     parser.add_argument("tfz_dir", help="Path to the tf_fuzz directory", type=Path)
     parser.add_argument(
-        "input_dir",
-        help="The directory containing test template files for the fuzzing tool",
+        "in_path",
+        help="a directory of *.test files, or the path to a single test file",
         type=Path,
     )
     parser.add_argument(
@@ -63,6 +63,13 @@
         "--seed", help="Random seed used during test generation", type=int
     )
 
+    parser.add_argument(
+        "-n",
+        help="Number of test cases to generate for each input test template",
+        type=int,
+        default=1,
+    )
+
     args = parser.parse_args()
 
     SEED: int = args.seed
@@ -72,7 +79,7 @@
     TFZ_DIR: Path = args.tfz_dir
     LIB_DIR: Path = TFZ_DIR / "lib"
     TFZ_EXECUTABLE: Path = args.tfz_dir / "bin" / "tfz"
-    INPUT_DIR: Path = args.input_dir
+    INPUT_PATH: Path = args.in_path
     TARGET_DIR: Path = args.target_dir
 
     if not TFZ_DIR.is_dir():
@@ -94,8 +101,16 @@
         parser.print_help()
         sys.exit(1)
 
-    if not INPUT_DIR.is_dir():
-        print(f"template_dir ({INPUT_DIR}) does not exist or is not a directory.")
+    if INPUT_PATH.is_dir():
+        file_iterator = enumerate(sorted(INPUT_PATH.glob("*.test")))
+
+    elif INPUT_PATH.is_file():
+        file_iterator = [(0, INPUT_PATH)]
+
+    else:
+        print(
+            f"path ({INPUT_PATH}) does not exist, is not a directory, or is not a file."
+        )
         parser.print_help()
         sys.exit(1)
 
@@ -117,58 +132,63 @@
     print(f"Using random seed: {SEED}")
     random.seed(SEED)
 
-    for i, test_input_path in enumerate(sorted(INPUT_DIR.glob("*.test"))):
-        # use different random seeds based on a known seed so the tests have
-        # different random values, but still can be ran deterministically
-        seed = random.randint(0, 0xFFFFFFFF)
+    print(f"Generating {args.n} cases for each test file")
 
+    for i, test_input_path in file_iterator:
         print(f"* Found test file {test_input_path}")
+        for j in range(0, args.n):
+            # use different random seeds based on a known seed so the tests have
+            # different random values, but still can be ran deterministically
+            seed = random.randint(0, 0xFFFFFFFF)
 
-        c_file_name: str = f"{test_input_path.stem}.c"
-        generated_test_path: Path = TARGET_DIR / f"{c_file_name}"
+            c_file_name: str = f"{test_input_path.stem}_{j}.c"
+            generated_test_path: Path = TARGET_DIR / f"{c_file_name}"
 
-        process = sp.run(
-            f"{TFZ_EXECUTABLE.absolute()} {test_input_path.absolute()} {generated_test_path.absolute()} {seed}",
-            shell=True,
-            text=True,
-            stderr=sp.STDOUT,
-            stdout=sp.PIPE,
-        )
+            process = sp.run(
+                f"{TFZ_EXECUTABLE.absolute()} {test_input_path.absolute()} {generated_test_path.absolute()} {seed}",
+                shell=True,
+                text=True,
+                stderr=sp.STDOUT,
+                stdout=sp.PIPE,
+            )
 
-        if process.returncode != 0:
-            print("  tfz invocation failed, skipping test")
-            print("  Command output: ")
-            for line in process.stdout.splitlines():
-                print(f"    {line}")
+            if process.returncode != 0:
+                print(f"  [{j + 1}/{args.n}] tfz invocation failed , skipping")
+                print("      Command output: ")
+                for line in process.stdout.splitlines():
+                    print(f"      {line}")
+                print()
+                continue
 
-            continue
+            name: str = (
+                f"TFM_FUZZ_TEST_{test_input_path.stem.upper().replace(' ','_')}_{j}"
+            )
 
-        name: str = f"TFM_FUZZ_TEST_{test_input_path.stem.upper().replace(' ','_')}"
+            # DESCRIPTION: use the purpose stored in the input test file.
+            # This is in the syntax `purpose to <str>;`.
+            description: str = ""
+            with open(test_input_path) as f:
+                _match = re.search(r"^purpose to\s+(.*);", f.read())
+                if _match:
+                    description = _match[1]
 
-        # DESCRIPTION: use the purpose stored in the input test file.
-        # This is in the syntax `purpose to <str>;`.
-        description: str = ""
-        with open(test_input_path) as f:
-            _match = re.search(r"^purpose to\s+(.*);", f.read())
-            if _match:
-                description = _match[1]
+            description = description.replace('"', '\\"')
 
-        description = description.replace('"', '\\"')
+            # FN_DEF / FN_NAME: the test function generated is always called
+            # test_thread. This is given a unique name, `test_<i>` when it is
+            # added to the suite.
 
-        # FN_DEF / FN_NAME: the test function generated is always called
-        # test_thread. This is given a unique name, `test_<i>` when it is
-        # added to the suite.
+            fn_name: str = f"test_tfz_generated_{i}_{j}"
 
-        fn_name: str = f"test_tfz_generated_{i}"
+            with open(generated_test_path, "r+") as generated_test_file:
 
-        with open(generated_test_path, "r+") as generated_test_file:
+                _contents = generated_test_file.read()
+                generated_test_file.seek(0)
+                generated_test_file.write(re.sub(r"test_thread", fn_name, _contents))
 
-            _contents = generated_test_file.read()
-            generated_test_file.seek(0)
-            generated_test_file.write(re.sub(r"test_thread", fn_name, _contents))
-
-        test: Test = Test(fn_name, c_file_name, name, description)
-        test_suite.tests.append(test)
+            test: Test = Test(fn_name, c_file_name, name, description)
+            test_suite.tests.append(test)
+            print(f"  [{j+1}/{args.n}] generated")
 
     for template_name in jinja_env.list_templates():
         template: Template = jinja_env.get_template(template_name)