Fix prolog.engine.check_answers

Variables not specified in the set of expected answers are now ignored.

1 files changed, 17 insertions, 12 deletions

diff --git a/prolog/engine.py b/prolog/engine.py
index eca0826..5258878 100644
--- a/prolog/engine.py
+++ b/prolog/engine.py
@@ -117,17 +117,26 @@ def pretty_vars(data):
 # within [timeout] seconds, and fails when it finds any other solution.
 def check_answers(engine, query, answers, timeout=10):
     seen = []
+    # Return false (only) if there is no expected answer where values of all
+    # specified variables match values in [answer].
+    def check_answer(answer):
+        if not answer:
+            return True
+        bindings, constraints = answer
+        for expected in answers:
+            if all(bindings.get(var) == val for var, val in expected.items()):
+                if expected not in seen:
+                    seen.append(expected)
+                return True
+        return False
+
     start = time.monotonic()
     try:
         # Run the query.
         reply, output = ask(engine, query, timeout)
         answer, error, more = process_answer(reply)
-        if answer:
-            bindings, constraints = answer
-            if bindings not in answers:
-                return False
-            if bindings not in seen:
-                seen.append(bindings)
+        if not check_answer(answer):
+            return False
 
         # Continue while there are more potential answers and time remaining.
         while more:
@@ -136,12 +145,8 @@ def check_answers(engine, query, answers, timeout=10):
                 break
             reply, output = next(engine, timeout=real_timeout)
             answer, error, more = process_answer(reply)
-            if answer:
-                bindings, constraints = answer
-                if bindings not in answers:
-                    return False
-                if bindings not in seen:
-                    seen.append(bindings)
+            if not check_answer(answer):
+                return False
     except:
         pass
     return len(seen) == len(answers)