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,33 +110,55 @@ usb_device.open()
 usb_device.reset()
 usb_device.close()
 
-print(f'### Starting USB capture on usbmon{bus_num}')
+def capture_test(cmd, capture_file, prompt):
-capture_pid = os.fork()
+    print(f'### Starting USB capture on usbmon{bus_num}')
-assert(capture_pid >= 0)
+    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(),
-if capture_pid == 0:
+                                       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)
+    # Wait 1 sec to settle (we can assume setpgrp happened)
-time.sleep(1)
+    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)
 
+    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
     r = os.waitpid(pid, os.WNOHANG)
@@ -146,29 +168,13 @@ def t_waitpid(pid, timeout):
 
     return r
 
-os.kill(capture_pid, signal.SIGTERM)
+capture_test([SRCDIR + '/tests/capture.py', test_dir + '/capture.png'],
-try:
+             'capture.pcapng',
-    r = t_waitpid(capture_pid, 2)
+             '### Capturing fingerprint, please swipe or press your finger on the reader')
-    # Kill if nothing died
-    if r[0] == 0:
-        os.kill(capture_pid, signal.SIGKILL)
-except ChildProcessError:
-    pass
 
-try:
+custom_script = os.path.join(test_dir, 'custom.py')
-    while True:
+if os.path.exists(custom_script):
-        r = t_waitpid(-capture_pid, timeout=2)
+    capture_test([custom_script], 'custom.pcapng',
-        # Kill the process group, if nothing died (and there are children)
+                 '### Running the custom fingerprint capture')
-        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()
 
 print(f"\nDone! Don't forget to add {test_name} to tests/meson.build")