Multithread the CLI tests

4 files changed, 115 insertions, 114 deletions

diff --git a/src/cli/encryption.cpp b/src/cli/encryption.cpp
index 4d85890b0..fa8de7cfd 100644
--- a/src/cli/encryption.cpp
+++ b/src/cli/encryption.cpp
@@ -18,7 +18,7 @@ namespace {
 
 auto VALID_MODES = std::map<std::string, std::string>{
          // Don't add algorithms here without extending tests
-         // in `src/scripts/cli_tests.py`
+         // in `src/scripts/test_cli_crypt.py`
          { "aes-128-cfb", "AES-128/CFB" },
          { "aes-192-cfb", "AES-192/CFB" },
          { "aes-256-cfb", "AES-256/CFB" },
diff --git a/src/scripts/ci_build.py b/src/scripts/ci_build.py
index 6a209faf1..ebbaf84f8 100755
--- a/src/scripts/ci_build.py
+++ b/src/scripts/ci_build.py
@@ -568,9 +568,11 @@ def main(args=None):
         if target in ['shared', 'coverage'] and options.os != 'windows':
             botan_exe = os.path.join(root_dir, 'botan-cli.exe' if options.os == 'windows' else 'botan')
 
-            test_scripts = ['cli_tests.py', 'test_cli.py']
+            test_scripts = ['test_cli.py', 'test_cli_crypt.py']
             for script in test_scripts:
-                cmds.append([py_interp, os.path.join(root_dir, 'src/scripts', script),
+                cmds.append([py_interp,
+                             os.path.join(root_dir, 'src/scripts', script),
+                             '--threads=%d' % (options.build_jobs),
                              botan_exe])
 
         python_tests = os.path.join(root_dir, 'src/scripts/test_python.py')
diff --git a/src/scripts/test_cli.py b/src/scripts/test_cli.py
index 661b2dbcb..fcf5d5018 100755
--- a/src/scripts/test_cli.py
+++ b/src/scripts/test_cli.py
@@ -18,6 +18,8 @@ import re
 import random
 import json
 import binascii
+import multiprocessing
+from multiprocessing.pool import ThreadPool
 
 # pylint: disable=global-statement,unused-argument
 
@@ -1105,6 +1107,46 @@ def cli_tls_client_hello_tests(_tmp_dir):
     output_hash = "8EBFC3205ACFA98461128FE5D081D19254237AF84F7DAF000A3C992C3CF6DE44"
     test_cli("hash", ["--no-fsname", "--algo=SHA-256", "-"], output_hash, output)
 
+def cli_speed_pk_tests(_tmp_dir):
+    msec = 1
+
+    pk_algos = ["ECDSA", "ECDH", "SM2", "ECKCDSA", "ECGDSA", "GOST-34.10",
+                "DH", "DSA", "ElGamal", "Ed25519", "Curve25519", "NEWHOPE", "McEliece",
+                "RSA", "XMSS"]
+
+    output = test_cli("speed", ["--msec=%d" % (msec)] + pk_algos, None).split('\n')
+
+    # ECDSA-secp256r1 106 keygen/sec; 9.35 ms/op 37489733 cycles/op (1 op in 9 ms)
+    format_re = re.compile(r'^.* [0-9]+ ([A-Za-z ]+)/sec; [0-9]+\.[0-9]+ ms/op .*\([0-9]+ (op|ops) in [0-9\.]+ ms\)')
+    for line in output:
+        if format_re.match(line) is None:
+            logging.error("Unexpected line %s", line)
+
+def cli_speed_pbkdf_tests(_tmp_dir):
+    msec = 1
+    pbkdf_ops = ['bcrypt', 'passhash9', 'argon2']
+
+    format_re = re.compile(r'^.* [0-9]+ /sec; [0-9]+\.[0-9]+ ms/op .*\([0-9]+ (op|ops) in [0-9]+(\.[0-9]+)? ms\)')
+    for op in pbkdf_ops:
+        output = test_cli("speed", ["--msec=%d" % (msec), op], None).split('\n')
+        for line in output:
+            if format_re.match(line) is None:
+                logging.error("Unexpected line %s", line)
+
+def cli_speed_math_tests(_tmp_dir):
+    msec = 1
+    # these all have a common output format
+    math_ops = ['mp_mul', 'mp_div', 'mp_div10', 'modexp', 'random_prime', 'inverse_mod',
+                'rfc3394', 'fpe_fe1', 'ecdsa_recovery', 'ecc_init', 'poly_dbl',
+                'bn_redc', 'nistp_redc', 'ecc_mult', 'ecc_ops', 'os2ecp', 'primality_test']
+
+    format_re = re.compile(r'^.* [0-9]+ /sec; [0-9]+\.[0-9]+ ms/op .*\([0-9]+ (op|ops) in [0-9]+(\.[0-9]+)? ms\)')
+    for op in math_ops:
+        output = test_cli("speed", ["--msec=%d" % (msec), op], None).split('\n')
+        for line in output:
+            if format_re.match(line) is None:
+                logging.error("Unexpected line %s", line)
+
 def cli_speed_tests(_tmp_dir):
     # pylint: disable=too-many-branches
 
@@ -1137,31 +1179,6 @@ def cli_speed_tests(_tmp_dir):
             if format_re_cipher.match(line) is None:
                 logging.error('Unexpected line %s', line)
 
-    pk_algos = ["ECDSA", "ECDH", "SM2", "ECKCDSA", "ECGDSA", "GOST-34.10",
-                "DH", "DSA", "ElGamal", "Ed25519", "Curve25519", "NEWHOPE", "McEliece",
-                "RSA", "XMSS"]
-
-    output = test_cli("speed", ["--msec=%d" % (msec)] + pk_algos, None).split('\n')
-
-    # ECDSA-secp256r1 106 keygen/sec; 9.35 ms/op 37489733 cycles/op (1 op in 9 ms)
-    format_re = re.compile(r'^.* [0-9]+ ([A-Za-z ]+)/sec; [0-9]+\.[0-9]+ ms/op .*\([0-9]+ (op|ops) in [0-9\.]+ ms\)')
-    for line in output:
-        if format_re.match(line) is None:
-            logging.error("Unexpected line %s", line)
-
-    # these all have a common output format
-    math_ops = ['mp_mul', 'mp_div', 'mp_div10', 'modexp', 'random_prime', 'inverse_mod',
-                'rfc3394', 'fpe_fe1', 'ecdsa_recovery', 'ecc_init', 'poly_dbl',
-                'bn_redc', 'nistp_redc', 'ecc_mult', 'ecc_ops', 'os2ecp', 'primality_test',
-                'bcrypt', 'passhash9', 'argon2']
-
-    format_re = re.compile(r'^.* [0-9]+ /sec; [0-9]+\.[0-9]+ ms/op .*\([0-9]+ (op|ops) in [0-9]+(\.[0-9]+)? ms\)')
-    for op in math_ops:
-        output = test_cli("speed", ["--msec=%d" % (msec), op], None).split('\n')
-        for line in output:
-            if format_re.match(line) is None:
-                logging.error("Unexpected line %s", line)
-
     output = test_cli("speed", ["--msec=%d" % (msec), "scrypt"], None).split('\n')
 
     format_re = re.compile(r'^scrypt-[0-9]+-[0-9]+-[0-9]+ \([0-9]+ MiB\) [0-9]+ /sec; [0-9]+\.[0-9]+ ms/op .*\([0-9]+ (op|ops) in [0-9\.]+ ms\)')
@@ -1196,8 +1213,20 @@ def cli_speed_tests(_tmp_dir):
             if field not in b:
                 logging.error('Missing field %s in JSON record %s' % (field, b))
 
+def run_test(fn_name, fn):
+    start = time.time()
+    tmp_dir = tempfile.mkdtemp(prefix='botan_cli_')
+    try:
+        fn(tmp_dir)
+    except Exception as e: # pylint: disable=broad-except
+        logging.error("Test %s threw exception: %s", fn_name, e)
+
+    shutil.rmtree(tmp_dir)
+    end = time.time()
+    logging.info("Ran %s in %.02f sec", fn_name, end-start)
 
 def main(args=None):
+    # pylint: disable=too-many-branches
     if args is None:
         args = sys.argv
 
@@ -1206,19 +1235,24 @@ def main(args=None):
 
     parser.add_option('--verbose', action='store_true', default=False)
     parser.add_option('--quiet', action='store_true', default=False)
+    parser.add_option('--threads', action='store', type='int', default=0)
 
     (options, args) = parser.parse_args(args)
 
     setup_logging(options)
 
     if len(args) < 2:
-        logging.error("Usage: ./cli_tests.py path_to_botan_cli [test_regex]")
+        logging.error("Usage: %s path_to_botan_cli [test_regex]", args[0])
         return 1
 
     if not os.access(args[1], os.X_OK):
         logging.error("Could not access/execute %s", args[1])
         return 2
 
+    threads = options.threads
+    if threads == 0:
+        threads = multiprocessing.cpu_count()
+
     global CLI_PATH
     CLI_PATH = args[1]
 
@@ -1230,9 +1264,14 @@ def main(args=None):
             logging.error("Invalid regex: %s", str(e))
             return 1
 
-    start_time = time.time()
-
+    # some of the slowest tests are grouped up front
     test_fns = [
+        cli_speed_tests,
+        cli_speed_pk_tests,
+        cli_speed_math_tests,
+        cli_speed_pbkdf_tests,
+        cli_xmss_sign_tests,
+
         cli_argon2_tests,
         cli_asn1_tests,
         cli_base32_tests,
@@ -1256,7 +1295,6 @@ def main(args=None):
         cli_hmac_tests,
         cli_is_prime_tests,
         cli_key_tests,
-        cli_xmss_sign_tests,
         cli_mod_inverse_tests,
         cli_pbkdf_tune_tests,
         cli_pk_encrypt_tests,
@@ -1265,7 +1303,6 @@ def main(args=None):
         cli_rng_tests,
         cli_roughtime_check_tests,
         cli_roughtime_tests,
-        cli_speed_tests,
         cli_timing_test_tests,
         cli_tls_ciphersuite_tests,
         cli_tls_client_hello_tests,
@@ -1278,25 +1315,26 @@ def main(args=None):
         cli_version_tests,
         ]
 
+    tests_to_run = []
     for fn in test_fns:
         fn_name = fn.__name__
 
-        if test_regex is not None:
-            if test_regex.search(fn_name) is None:
-                continue
+        if test_regex is None or test_regex.search(fn_name) is not None:
+            tests_to_run.append((fn_name, fn))
 
-        logging.info("Running %s" % (fn_name))
+    start_time = time.time()
 
-        start = time.time()
-        tmp_dir = tempfile.mkdtemp(prefix='botan_cli_')
-        try:
-            fn(tmp_dir)
-        except Exception as e: # pylint: disable=broad-except
-            logging.error("Test %s threw exception: %s", fn_name, e)
+    if threads > 1:
+        pool = ThreadPool(processes=threads)
+        results = []
+        for test in tests_to_run:
+            results.append(pool.apply_async(run_test, test))
 
-        shutil.rmtree(tmp_dir)
-        end = time.time()
-        logging.debug("Ran %s in %.02f sec", fn_name, end-start)
+        for result in results:
+            result.get()
+    else:
+        for test in tests_to_run:
+            run_test(test[0], test[1])
 
     end_time = time.time()
 
diff --git a/src/scripts/cli_tests.py b/src/scripts/test_cli_crypt.py
index b4e973c2b..780c14314 100755
--- a/src/scripts/cli_tests.py
+++ b/src/scripts/test_cli_crypt.py
@@ -13,23 +13,18 @@ import vecparser
 
 cli_binary = ""
 
-SUPPORTED_ALGORITHMS = [
-    'AES-128/CFB',
-    'AES-192/CFB',
-    'AES-256/CFB',
-    'AES-128/GCM',
-    'AES-192/GCM',
-    'AES-256/GCM',
-    'AES-128/OCB',
-    'AES-128/XTS',
-    'AES-256/XTS',
-    'ChaCha20Poly1305',
-]
-
-def append_ordered(base, additional_elements):
-    for key in additional_elements:
-        value = additional_elements[key]
-        base[key] = value
+SUPPORTED_ALGORITHMS = {
+    "AES-128/CFB": "aes-128-cfb",
+    "AES-192/CFB": "aes-192-cfb",
+    "AES-256/CFB": "aes-256-cfb",
+    "AES-128/GCM": "aes-128-gcm",
+    "AES-192/GCM": "aes-192-gcm",
+    "AES-256/GCM": "aes-256-gcm",
+    "AES-128/OCB": "aes-128-ocb",
+    "AES-128/XTS": "aes-128-xts",
+    "AES-256/XTS": "aes-256-xts",
+    "ChaCha20Poly1305": "chacha20poly1305",
+}
 
 class TestSequence(unittest.TestCase):
     pass
@@ -44,32 +39,8 @@ def create_test(data):
         algorithm = data['Algorithm']
         direction = data['Direction']
 
-        # CFB
-        if algorithm == "AES-128/CFB":
-            mode = "aes-128-cfb"
-        elif algorithm == "AES-192/CFB":
-            mode = "aes-192-cfb"
-        elif algorithm == "AES-256/CFB":
-            mode = "aes-256-cfb"
-        # GCM
-        elif algorithm == "AES-128/GCM":
-            mode = "aes-128-gcm"
-        elif algorithm == "AES-192/GCM":
-            mode = "aes-192-gcm"
-        elif algorithm == "AES-256/GCM":
-            mode = "aes-256-gcm"
-        # OCB
-        elif algorithm == "AES-128/OCB":
-            mode = "aes-128-ocb"
-        # XTS
-        elif algorithm == "AES-128/XTS":
-            mode = "aes-128-xts"
-        elif algorithm == "AES-256/XTS":
-            mode = "aes-256-xts"
-        # ChaCha20Poly1305
-        elif algorithm == "ChaCha20Poly1305":
-            mode = "chacha20poly1305"
-        else:
+        mode = SUPPORTED_ALGORITHMS.get(algorithm)
+        if mode is None:
             raise Exception("Unknown algorithm: '" + algorithm + "'")
 
         cmd = [
@@ -81,15 +52,12 @@ def create_test(data):
             "--key=%s" % key]
         if direction == "decrypt":
             cmd += ['--decrypt']
-        # out_raw = subprocess.check_output(cmd)
 
         if direction == "decrypt":
             invalue = ciphertext
         else:
             invalue = plaintext
 
-        #print(cmd)
-
         p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stdin=subprocess.PIPE, stderr=subprocess.PIPE)
         out_raw = p.communicate(input=binascii.unhexlify(invalue))[0]
         out = binascii.hexlify(out_raw).decode("UTF-8").lower()
@@ -132,36 +100,29 @@ if __name__ == '__main__':
     parser.add_argument('cli_binary',
                         help='path to the botan cli binary')
     parser.add_argument('--max-tests', type=int, default=20)
+    parser.add_argument('--threads', type=int, default=0)
     parser.add_argument('unittest_args', nargs="*")
     args = parser.parse_args()
 
     cli_binary = args.cli_binary
     max_tests = args.max_tests
 
-    vecfile_cfb = vecparser.VecDocument(os.path.join('src', 'tests', 'data', 'modes', 'cfb.vec'))
-    vecfile_gcm = vecparser.VecDocument(os.path.join('src', 'tests', 'data', 'aead', 'gcm.vec'))
-    vecfile_ocb = vecparser.VecDocument(os.path.join('src', 'tests', 'data', 'aead', 'ocb.vec'))
-    vecfile_xts = vecparser.VecDocument(os.path.join('src', 'tests', 'data', 'modes', 'xts.vec'))
-    vecfile_chacha20poly1305 = vecparser.VecDocument(os.path.join('src', 'tests', 'data', 'aead', 'chacha20poly1305.vec'))
-    #data = vecfile.get_data()
-    #for algo in data:
-    #    print(algo)
-    #    i = 0
-    #    for testcase in data[algo]:
-    #        i += 1
-    #        print(str(i) + ":", testcase)
+    test_data_dir = os.path.join('src', 'tests', 'data')
+
+    mode_test_data = [os.path.join(test_data_dir, 'modes', 'cfb.vec'),
+                      os.path.join(test_data_dir, 'aead', 'gcm.vec'),
+                      os.path.join(test_data_dir, 'aead', 'ocb.vec'),
+                      os.path.join(test_data_dir, 'modes', 'xts.vec'),
+                      os.path.join(test_data_dir, 'aead', 'chacha20poly1305.vec')]
 
     testdata = OrderedDict()
-    append_ordered(testdata, get_testdata(vecfile_cfb.get_data(), max_tests))
-    append_ordered(testdata, get_testdata(vecfile_gcm.get_data(), max_tests))
-    append_ordered(testdata, get_testdata(vecfile_ocb.get_data(), max_tests))
-    append_ordered(testdata, get_testdata(vecfile_xts.get_data(), max_tests))
-    append_ordered(testdata, get_testdata(vecfile_chacha20poly1305.get_data(), max_tests))
-
-    #for testname in testdata:
-    #    print(testname)
-    #    for key in testdata[testname]:
-    #        print("    " + key + ": " + testdata[testname][key])
+
+    for f in mode_test_data:
+        vecfile = vecparser.VecDocument(f)
+        vecdata = get_testdata(vecfile.get_data(), max_tests)
+        for key in vecdata:
+            testdata[key] = vecdata[key]
+
     for testname in testdata:
         test_method = create_test(testdata[testname])
         test_method.__name__ = 'test_%s' % testname