From f331ee415534f15e52161ec909a22136e20adce7 Mon Sep 17 00:00:00 2001
From: Hanno Becker <beckphan@amazon.co.uk>
Date: Wed, 1 Jan 2025 22:07:56 +0000
Subject: [PATCH] Tests: Reduce amount of nesting in benchmarking JSON output

Signed-off-by: Hanno Becker <beckphan@amazon.co.uk>
---
 scripts/lib/mlkem_test.py | 53 +++++++++++++++++++--------------------
 1 file changed, 26 insertions(+), 27 deletions(-)

diff --git a/scripts/lib/mlkem_test.py b/scripts/lib/mlkem_test.py
index 5f5bd8c64..96b2ebf49 100644
--- a/scripts/lib/mlkem_test.py
+++ b/scripts/lib/mlkem_test.py
@@ -277,34 +277,33 @@ def bench(self):
 
         # NOTE: There will only be one items in resultss, as we haven't yet decided how to write both opt/no-opt benchmark results
         for k, results in resultss.items():
-            if results is not None and output is not None and components is False:
-                import json
-
-                with open(output, "w") as f:
-                    v = []
-                    for scheme in results:
-                        schemeStr = str(scheme)
-                        r = results[scheme]
-
-                        # The first 3 lines of the output are expected to be
-                        # keypair cycles=X
-                        # encaps cycles=X
-                        # decaps cycles=X
-
-                        lines = [line for line in r.splitlines() if "=" in line]
-
-                        d = {
-                            k.strip(): int(v) for k, v in (l.split("=") for l in lines)
+            if not (results is not None and output is not None and components is False):
+                continue
+
+            v = []
+            for scheme in results:
+                schemeStr = str(scheme)
+                r = results[scheme]
+
+                # The first 3 lines of the output are expected to be
+                # keypair cycles=X
+                # encaps cycles=X
+                # decaps cycles=X
+
+                lines = [line for line in r.splitlines() if "=" in line]
+
+                d = {k.strip(): int(v) for k, v in (l.split("=") for l in lines)}
+                for primitive in ["keypair", "encaps", "decaps"]:
+                    v.append(
+                        {
+                            "name": f"{schemeStr} {primitive}",
+                            "unit": "cycles",
+                            "value": d[f"{primitive} cycles"],
                         }
-                        for primitive in ["keypair", "encaps", "decaps"]:
-                            v.append(
-                                {
-                                    "name": f"{schemeStr} {primitive}",
-                                    "unit": "cycles",
-                                    "value": d[f"{primitive} cycles"],
-                                }
-                            )
-                    f.write(json.dumps(v))
+                    )
+
+            with open(output, "w") as f:
+                f.write(json.dumps(v))
 
     def all(self):
         func = self.args.func