diff --git a/test/pytest/generate_ci_yaml.py b/test/pytest/generate_ci_yaml.py
index 7a495267ab..f62d752ade 100644
--- a/test/pytest/generate_ci_yaml.py
+++ b/test/pytest/generate_ci_yaml.py
@@ -1,6 +1,6 @@
-import glob
 import itertools
 import os
+from pathlib import Path
 
 import yaml
 
@@ -9,6 +9,7 @@
 in the pytests directory to parallelise the CI jobs.
 '''
 
+
 template = """
 pytest.{}:
   extends: .pytest
@@ -19,6 +20,14 @@
 
 n_test_files_per_yml = int(os.environ.get('N_TESTS_PER_YAML', 4))
 
+BLACKLIST = {'test_reduction'}
+
+
+def path_to_name(test_path):
+    path = Path(test_path)
+    name = path.stem.replace('test_', '')
+    return name
+
 
 def batched(iterable, chunk_size):
     iterator = iter(iterable)
@@ -32,41 +41,32 @@ def uses_example_model(test_filename):
         return 'example-models' in content
 
 
-yml = None
-tests = glob.glob('test_*.py')
-for test_batch in batched(tests, n_test_files_per_yml):
-    name = '+'.join([test.replace('test_', '').replace('.py', '') for test in test_batch])
-    test_files = ' '.join(list(test_batch))
-    uses_example_models = int(any([uses_example_model(test) for test in test_batch]))
-
-    new_yml = yaml.safe_load(template.format(name, test_files, uses_example_models))
-    if yml is None:
-        yml = new_yml
-    else:
-        yml.update(new_yml)
-
-# hls4ml Optimization API
-tests = glob.glob('test_optimization/test_*.py')
-for test in tests:
-    name = test.replace('test_optimization/', '').replace('test_', '').replace('.py', '')
-    new_yml = yaml.safe_load(template.format(name, f'test_optimization/test_{name}.py', int(uses_example_model(test))))
-    if yml is None:
-        yml = new_yml
-    else:
-        yml.update(new_yml)
-
-tests = glob.glob('test_optimization/test_keras/test_*.py')
-for test in tests:
-    # For now, skip Keras Surgeon [conflicting versions]
-    if 'test_reduction' not in test:
-        name = test.replace('test_optimization/test_keras/', '').replace('test_', '').replace('.py', '')
-        new_yml = yaml.safe_load(
-            template.format(name, f'test_optimization/test_keras/test_{name}.py', int(uses_example_model(test)))
-        )
+def generate_test_yaml(test_root='.'):
+    test_root = Path(test_root)
+    test_paths = [path for path in test_root.glob('**/test_*.py') if path.stem not in BLACKLIST]
+    for path in test_paths:
+        print(path.name)
+    need_example_models = [uses_example_model(path) for path in test_paths]
+
+    idxs = list(range(len(need_example_models)))
+    idxs = sorted(idxs, key=lambda i: f'{need_example_models[i]}_{path_to_name(test_paths[i])}')
+
+    yml = None
+    for batch_idxs in batched(idxs, n_test_files_per_yml):
+        batch_paths: list[Path] = [test_paths[i] for i in batch_idxs]
+        names = [path_to_name(path) for path in batch_paths]
+        name = '+'.join(names)
+        test_files = ' '.join([str(path.relative_to(test_root)) for path in batch_paths])
+        batch_need_example_model = int(any([need_example_models[i] for i in batch_idxs]))
+        diff_yml = yaml.safe_load(template.format(name, test_files, batch_need_example_model))
         if yml is None:
-            yml = new_yml
+            yml = diff_yml
         else:
-            yml.update(new_yml)
+            yml.update(diff_yml)
+    return yml
+
 
-yamlfile = open('pytests.yml', 'w')
-yaml.safe_dump(yml, yamlfile)
+if __name__ == '__main__':
+    yml = generate_test_yaml(Path(__file__).parent)
+    with open('pytests.yml', 'w') as yamlfile:
+        yaml.safe_dump(yml, yamlfile)