Update mock_api.py

neulab · Dec 8, 2023 · bc5de3d · bc5de3d
1 parent 74213f8
commit bc5de3d
Showing 1 changed file with 75 additions and 15 deletions.
diff --git a/test_helpers/mock_api.py b/test_helpers/mock_api.py
@@ -100,15 +100,39 @@ def __init__(self, length: int = 4) -> None:
         self.current_index = 0
         mock_completion_1 = MockCompletion()
         mock_completion_1.choices = [
-            {"message": {"content": '{"input": "1", "explanation": "x", "output": "a"}'}},
-            {"message": {"content": '{"input": "1", "explanation": "x", "output": "b"}'}},
-            {"message": {"content": '{"input": "1", "explanation": "x", "output": "a"}'}},
+            {
+                "message": {
+                    "content": '{"input": "1", "explanation": "x", "output": "a"}'
+                }
+            },
+            {
+                "message": {
+                    "content": '{"input": "1", "explanation": "x", "output": "b"}'
+                }
+            },
+            {
+                "message": {
+                    "content": '{"input": "1", "explanation": "x", "output": "a"}'
+                }
+            },
         ]
         mock_completion_2 = MockCompletion()
         mock_completion_2.choices = [
-            {"message": {"content": '{"input": "1", "explanation": "x", "output": "c"}'}},
-            {"message": {"content": '{"input": "2", "explanation": "x", "output": "a"}'}},
-            {"message": {"content": '{"input": "2", "explanation": "x", "output": "b"}'}},
+            {
+                "message": {
+                    "content": '{"input": "1", "explanation": "x", "output": "c"}'
+                }
+            },
+            {
+                "message": {
+                    "content": '{"input": "2", "explanation": "x", "output": "a"}'
+                }
+            },
+            {
+                "message": {
+                    "content": '{"input": "2", "explanation": "x", "output": "b"}'
+                }
+            },
         ]
         self.mock_completions.append(
             [
@@ -118,24 +142,60 @@ def __init__(self, length: int = 4) -> None:
         )
         mock_completion_3 = MockCompletion()
         mock_completion_3.choices = [
-            {"message": {"content": '{"input": "3", "explanation": "x", "output": "a"}'}},
-            {"message": {"content": '{"input": "3", "explanation": "x", "output": "a"}'}},
-            {"message": {"content": '{"input": "3", "explanation": "x", "output": "b"}'}},
+            {
+                "message": {
+                    "content": '{"input": "3", "explanation": "x", "output": "a"}'
+                }
+            },
+            {
+                "message": {
+                    "content": '{"input": "3", "explanation": "x", "output": "a"}'
+                }
+            },
+            {
+                "message": {
+                    "content": '{"input": "3", "explanation": "x", "output": "b"}'
+                }
+            },
         ]
         self.mock_completions.append([mock_completion_3])
 
         mock_completion_4 = MockCompletion()
         mock_completion_4.choices = [
-            {"message": {"content": '{"input": "1", "explanation": "x", "output": "b"}'}},
-            {"message": {"content": '{"input": "1", "explanation": "x", "output": "b"}'}},
-            {"message": {"content": '{"input": "1", "explanation": "x", "output": "b"}'}},
+            {
+                "message": {
+                    "content": '{"input": "1", "explanation": "x", "output": "b"}'
+                }
+            },
+            {
+                "message": {
+                    "content": '{"input": "1", "explanation": "x", "output": "b"}'
+                }
+            },
+            {
+                "message": {
+                    "content": '{"input": "1", "explanation": "x", "output": "b"}'
+                }
+            },
         ]
         self.mock_completions.append([mock_completion_4])
         mock_completion_5 = MockCompletion()
         mock_completion_5.choices = [
-            {"message": {"content": '{"input": "4", "explanation": "x", "output": "c"}'}},
-            {"message": {"content": '{"input": "4", "explanation": "x", "output": "c"}'}},
-            {"message": {"content": '{"input": "5", "explanation": "x", "output": "a"}'}},
+            {
+                "message": {
+                    "content": '{"input": "4", "explanation": "x", "output": "c"}'
+                }
+            },
+            {
+                "message": {
+                    "content": '{"input": "4", "explanation": "x", "output": "c"}'
+                }
+            },
+            {
+                "message": {
+                    "content": '{"input": "5", "explanation": "x", "output": "a"}'
+                }
+            },
         ]
         self.mock_completions.append([mock_completion_5])
         if length == 5: