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,11 +110,13 @@ usb_device.open()
 usb_device.reset()
 usb_device.close()
 
+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')
+    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]
@@ -123,29 +125,14 @@ if capture_pid == 0:
     # 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')
+    print(prompt)
-cmd = ['python3', SRCDIR + '/tests/capture.py', test_dir + '/capture.png']
+    with subprocess.Popen([sys.executable] + cmd) as capture_process:
-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"
-
-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)
 
-def t_waitpid(pid, timeout):
-    timeout = time.time() + timeout
-    r = os.waitpid(pid, os.WNOHANG)
-    while timeout > time.time() and r[0] == 0:
-        time.sleep(0.1)
-        r = os.waitpid(pid, os.WNOHANG)
-
-    return r
-
     os.kill(capture_pid, signal.SIGTERM)
     try:
         r = t_waitpid(capture_pid, 2)
@@ -166,9 +153,28 @@ except ChildProcessError:
 
     # 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}',
+    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
+    r = os.waitpid(pid, os.WNOHANG)
+    while timeout > time.time() and r[0] == 0:
+        time.sleep(0.1)
+        r = os.waitpid(pid, os.WNOHANG)
+
+    return r
+
+capture_test([SRCDIR + '/tests/capture.py', test_dir + '/capture.png'],
+             'capture.pcapng',
+             '### Capturing fingerprint, please swipe or press your finger on the reader')
+
+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")