Fix test function call in monkey.edits

1 files changed, 2 insertions, 1 deletions

diff --git a/monkey/edits.py b/monkey/edits.py
index 82e5066..dfdf164 100644
--- a/monkey/edits.py
+++ b/monkey/edits.py
@@ -57,7 +57,8 @@ def get_edits_from_trace(trace, test, id):
             if action.type == 'prolog_solve':
                 queries.add(action.query.rstrip(' .'))
             elif action.type == 'test':
-                correct = test(code)
+                passed, total = test(code)
+                correct = passed == total
                 submissions.add((code, correct))
                 if correct:
                     # Ignore actions after the first correct version.