update e2e test

HazyResearch · Dec 8, 2018 · c731123 · c731123
1 parent b5d06d7
commit c731123
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/e2e/test_e2e.py b/tests/e2e/test_e2e.py
@@ -370,9 +370,9 @@ def test_e2e(caplog):
     assert session.query(PartVolt).count() == 4282
 
     # Grab candidate lists
-    train_cands = candidate_extractor.get_candidates(split=0)
-    dev_cands = candidate_extractor.get_candidates(split=1)
-    test_cands = candidate_extractor.get_candidates(split=2)
+    train_cands = candidate_extractor.get_candidates(split=0, sort=True)
+    dev_cands = candidate_extractor.get_candidates(split=1, sort=True)
+    test_cands = candidate_extractor.get_candidates(split=2, sort=True)
     assert len(train_cands) == 2
     assert len(train_cands[0]) == 3684
     assert (