From 4f1fccd6a0a8106e4362ff80854f060f4fb967d3 Mon Sep 17 00:00:00 2001
From: djl11 <daniel.lenton11@imperial.ac.uk>
Date: Tue, 27 Aug 2024 22:12:14 -0700
Subject: [PATCH] added simple hack for the code blocks which assume prior
 imports of unify, as per earlier code blocks.

---
 tests/helpers.py | 41 ++++++++++++++++++++++++-----------------
 1 file changed, 24 insertions(+), 17 deletions(-)

diff --git a/tests/helpers.py b/tests/helpers.py
index 6ad63f899..54fca7392 100644
--- a/tests/helpers.py
+++ b/tests/helpers.py
@@ -26,21 +26,6 @@ def _replace_api_key_placeholders(example: str) -> str:
     return example
 
 
-def _test_examples(examples: List[str], executable: callable) -> Tuple[Dict[str, Union[True, str]], bool]:
-    all_passed = True
-    results = dict()
-    for example in examples:
-        example_w_key = _replace_api_key_placeholders(example)
-        try:
-            executable(example_w_key)
-            val = True
-        except Exception as e:
-            val = str(e)
-            all_passed = False
-        results[example] = val
-    return results, all_passed
-
-
 def _capture_exec_output(code_str: str) -> Any:
     with io.StringIO() as buf, redirect_stdout(buf):
         exec(code_str)
@@ -59,14 +44,36 @@ def _test_python_fn(str_in: str) -> None:
         ret = _capture_exec_output(str_in)
         if "import requests" in str_in:
             _check_rest_api_failures(ret)
-    return _test_examples(examples, _test_python_fn)
+    all_passed = True
+    results = dict()
+    for example in examples:
+        example_w_key = "import unify\n" + _replace_api_key_placeholders(example)
+        try:
+            _test_python_fn(example_w_key)
+            val = True
+        except Exception as e:
+            val = str(e)
+            all_passed = False
+        results[example] = val
+    return results, all_passed
 
 
 def _test_shell_examples(examples: List[str]) -> Tuple[Dict[str, Union[True, str]], bool]:
     def _test_shell_fn(str_in: str) -> None:
         ret = os.popen(str_in).read()
         _check_rest_api_failures(ret)
-    return _test_examples(examples, _test_shell_fn)
+    all_passed = True
+    results = dict()
+    for example in examples:
+        example_w_key = _replace_api_key_placeholders(example)
+        try:
+            _test_shell_fn(example_w_key)
+            val = True
+        except Exception as e:
+            val = str(e)
+            all_passed = False
+        results[example] = val
+    return results, all_passed
 
 
 def get_mdx_filepaths() -> List[str]: