From d613dc6a9a4fdc39899eebbaa33f24dc12820f76 Mon Sep 17 00:00:00 2001
From: Michael Fitzmaurice <Michael.Fitzmaurice@arup.com>
Date: Wed, 12 Oct 2022 22:30:48 +0100
Subject: [PATCH] Use multiprocessing in the notebooks smoketest

---
 .gitignore                        |  1 +
 notebooks/notebook_smoke_tests.py | 24 ++++++++++++++++--------
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/.gitignore b/.gitignore
index eb5f40aa..c83266cb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -133,3 +133,4 @@ scripts/secrets-mgr-exploration.py
 example_data/pt2matsim_network/genet_output/*
 genet_output/*
 example_data/output_*/*
+example_data/api_requests_send.json
diff --git a/notebooks/notebook_smoke_tests.py b/notebooks/notebook_smoke_tests.py
index f1a027c9..1bcd563e 100644
--- a/notebooks/notebook_smoke_tests.py
+++ b/notebooks/notebook_smoke_tests.py
@@ -1,8 +1,11 @@
 import argparse
 import glob
+import multiprocessing
 import subprocess
 import sys
 from datetime import datetime
+from itertools import repeat
+from multiprocessing import Manager
 from pprint import pprint
 
 from colorama import Fore, Style
@@ -64,7 +67,7 @@ def install_jupyter_kernel(kernel_name):
     return run_shell_command(kernel_install_cmd)
 
 
-def execute_notebook(notebook_path):
+def execute_notebook(notebook_path, notebook_results_dict):
     print("Executing notebook '{}{}{}'...".format(Fore.YELLOW, notebook_path, Style.RESET_ALL))
     cmd = [
         'jupyter', 'nbconvert',
@@ -72,7 +75,8 @@ def execute_notebook(notebook_path):
         '--execute', '"{}"'.format(notebook_path),
         '--output-dir=/tmp'
     ]
-    return run_shell_command(cmd)
+    return_code, cmd, run_time = run_shell_command(cmd)
+    notebook_results_dict[notebook_path] = (return_code, run_time)
 
 
 def run_shell_command(shell_cmd):
@@ -80,7 +84,13 @@ def run_shell_command(shell_cmd):
     start_time = datetime.now()
     rc = subprocess.call(' '.join(shell_cmd), shell=True)
     running_time = datetime.now() - start_time
-    print("{}Shell process return value was {}{}{}".format(Style.RESET_ALL, Fore.YELLOW, rc, Style.RESET_ALL))
+    print("{}Shell process return value for {}{}{} was {}{}{}".format(Style.RESET_ALL,
+                                                                      Fore.YELLOW,
+                                                                      shell_cmd,
+                                                                      Style.RESET_ALL,
+                                                                      Fore.YELLOW,
+                                                                      rc,
+                                                                      Style.RESET_ALL))
     return rc, ' '.join(shell_cmd), running_time
 
 
@@ -154,11 +164,9 @@ def print_summary(notebook_results_dict):
         print("{}Warning: Jupyter kernel installation shell command did not exit normally"
               " - this may cause problems later{}".format(Fore.RED, Style.RESET_ALL))
 
-    notebook_results = {}
-    for notebook in notebooks:
-        print('------------------------------------------------------')
-        return_code, cmd, run_time = execute_notebook(notebook)
-        notebook_results[notebook] = (return_code, run_time)
+    notebook_results = Manager().dict()
+    with multiprocessing.Pool() as pool:
+        pool.starmap(execute_notebook, zip(notebooks, repeat(notebook_results)))
 
     print('------------------------------------------------------')
     print("\nFinished the smoke test in {}{}{}".format(Fore.YELLOW, datetime.now() - start, Style.RESET_ALL))