fix test

terrierteam · Dec 2, 2024 · 069a28c · 069a28c
1 parent 7594ed3
commit 069a28c
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/tests/test_models.py b/tests/test_models.py
@@ -156,10 +156,8 @@ def test_e5(self):
         from pyterrier_dr import E5
         testmodel = E5.base()
         self._base_test(testmodel)
-        self.assertEqual(
-            True,
-            testmodel([{'qid': 'q1', 'query' : 'chemical reactions', 'docno' : 'd2', 'text' : 'professor proton mixed the chemical'}])[0]['score'] >0,
-        )
+        inp = pd.DataFrame([{'qid': 'q1', 'query' : 'chemical reactions', 'docno' : 'd2', 'text' : 'professor proton mixed the chemical'}])
+        self.assertTrue(testmodel(inp).iloc[0]['score'] > 0)
 
     def test_tasb(self):
         from pyterrier_dr import TasB