Nit fix for the ignored test (#407)

google · Dec 12, 2023 · 8905c0a · 8905c0a
1 parent 058c5f6
commit 8905c0a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/scripts/tests/test_train.py b/scripts/tests/test_train.py
@@ -118,7 +118,7 @@ def test_standard_setup(self) -> None:
     else:
       raise ValueError('val_dataset is not an instance of Dataset.')
 
-  def teset_no_val(self) -> None:
+  def test_no_val(self) -> None:
     train_data_path = tempfile.NamedTemporaryFile().name
     with open(train_data_path, 'w') as f:
       f.write(('1\tfoo\tbar\n'
@@ -128,7 +128,7 @@ def teset_no_val(self) -> None:
                '-1\tbaz\tqux\n'))
     train_dataset, features, val_dataset = train.preprocess(train_data_path, 1)
     self.assertEqual(features, ['foo', 'bar', 'baz'])
-    self.assertEqual(train_dataset.Y.tolist(), [True, False, True, True, False])
+    self.assertEqual(train_dataset.Y.tolist(), [1, -1, 1, 1, -1])
     self.assertEqual(train_dataset.X_rows.tolist(), [0, 0, 1, 2, 2, 2, 3, 3, 4])
     self.assertEqual(train_dataset.X_cols.tolist(), [0, 1, 0, 0, 1, 2, 1, 0, 2])
     self.assertIsNone(val_dataset)