tests/create-driver-tests.py: Allow to create both capture and custom tests

A device may provide both tests, so let's make this easier to achieve

tests/create-driver-test.py.in

@@ -110,32 +110,54 @@ usb_device.open()
 usb_device.reset()
 usb_device.close()
 
-print(f'### Starting USB capture on usbmon{bus_num}')
-capture_pid = os.fork()
-assert(capture_pid >= 0)
+def capture_test(cmd, capture_file, prompt):
+    print(f'### Starting USB capture on usbmon{bus_num}')
+    capture_pid = os.fork()
+    assert(capture_pid >= 0)
 
-unfiltered_cap_path = os.path.join(tempfile.gettempdir(), 'capture-unfiltered.pcapng')
-if capture_pid == 0:
-    os.setpgrp()
-    args = ['tshark', '-q', '-i', f'usbmon{bus_num}', '-w', unfiltered_cap_path]
-    os.execv(tshark, args)
+    unfiltered_cap_path = os.path.join(tempfile.gettempdir(),
+                                       f'{capture_file}-unfiltered.pcapng')
+    if capture_pid == 0:
+        os.setpgrp()
+        args = ['tshark', '-q', '-i', f'usbmon{bus_num}', '-w', unfiltered_cap_path]
+        os.execv(tshark, args)
 
-# Wait 1 sec to settle (we can assume setpgrp happened)
-time.sleep(1)
+    # Wait 1 sec to settle (we can assume setpgrp happened)
+    time.sleep(1)
 
-print('### Capturing fingerprint, please swipe or press your finger on the reader')
-cmd = ['python3', SRCDIR + '/tests/capture.py', test_dir + '/capture.png']
-capture_file = 'capture.pcapng' # capture for "capture" test
-if os.path.exists(os.path.join(test_dir, "custom.py")):
-    cmd = ['python3', os.path.join(test_dir, "custom.py")]
-    capture_file = "custom.pcapng"
+    print(prompt)
+    with subprocess.Popen([sys.executable] + cmd) as capture_process:
+        capture_process.wait()
+        if capture_process.returncode != 0:
+            print('Failed to capture fingerprint')
+            os.killpg(capture_pid, signal.SIGKILL)
+            sys.exit(1)
 
-with subprocess.Popen(cmd) as capture_process:
-    capture_process.wait()
-    if capture_process.returncode != 0:
-        print('Failed to capture fingerprint')
-        os.killpg(capture_pid, signal.SIGKILL)
-        sys.exit(1)
+    os.kill(capture_pid, signal.SIGTERM)
+    try:
+        r = t_waitpid(capture_pid, 2)
+        # Kill if nothing died
+        if r[0] == 0:
+            os.kill(capture_pid, signal.SIGKILL)
+    except ChildProcessError:
+        pass
+
+    try:
+        while True:
+            r = t_waitpid(-capture_pid, timeout=2)
+            # Kill the process group, if nothing died (and there are children)
+            if r[0] == 0:
+                os.killpg(capture_pid, signal.SIGKILL)
+    except ChildProcessError:
+        pass
+
+    # Filter the capture
+    print(f'\n### Saving USB capture as test case {test_name}')
+    args = ['tshark', '-r', unfiltered_cap_path,
+            '-Y', f'usb.bus_id == {bus_num} and usb.device_address == {device_num}',
+            '-w', os.path.join(test_dir, capture_file)]
+    with subprocess.Popen(args, stderr=subprocess.DEVNULL) as filter_process:
+        filter_process.wait()
 
 def t_waitpid(pid, timeout):
     timeout = time.time() + timeout
@@ -146,29 +168,13 @@ def t_waitpid(pid, timeout):
 
     return r
 
-os.kill(capture_pid, signal.SIGTERM)
-try:
-    r = t_waitpid(capture_pid, 2)
-    # Kill if nothing died
-    if r[0] == 0:
-        os.kill(capture_pid, signal.SIGKILL)
-except ChildProcessError:
-    pass
+capture_test([SRCDIR + '/tests/capture.py', test_dir + '/capture.png'],
+             'capture.pcapng',
+             '### Capturing fingerprint, please swipe or press your finger on the reader')
 
-try:
-    while True:
-        r = t_waitpid(-capture_pid, timeout=2)
-        # Kill the process group, if nothing died (and there are children)
-        if r[0] == 0:
-            os.killpg(capture_pid, signal.SIGKILL)
-except ChildProcessError:
-    pass
-
-# Filter the capture
-print(f'\n### Saving USB capture as test case {test_name}')
-args = ['tshark', '-r', unfiltered_cap_path, '-Y', f'usb.bus_id == {bus_num} and usb.device_address == {device_num}',
-        '-w', os.path.join(test_dir, capture_file)]
-with subprocess.Popen(args, stderr=subprocess.DEVNULL) as filter_process:
-    filter_process.wait()
+custom_script = os.path.join(test_dir, 'custom.py')
+if os.path.exists(custom_script):
+    capture_test([custom_script], 'custom.pcapng',
+                 '### Running the custom fingerprint capture')
 
 print(f"\nDone! Don't forget to add {test_name} to tests/meson.build")