MilesCranmer commited on
Commit
eae8f9c
1 Parent(s): ecefe62

Add test for yes/no function

Browse files
Files changed (1) hide show
  1. test/test.py +10 -1
test/test.py CHANGED
@@ -1,7 +1,8 @@
1
  import unittest
 
2
  import numpy as np
3
  from pysr import pysr, get_hof, best, best_tex, best_callable, best_row
4
- from pysr.sr import run_feature_selection, _handle_feature_selection
5
  import sympy
6
  from sympy import lambdify
7
  import pandas as pd
@@ -147,3 +148,11 @@ class TestFeatureSelection(unittest.TestCase):
147
  np.testing.assert_array_equal(
148
  np.sort(selected_X, axis=1), np.sort(X[:, [2, 3]], axis=1)
149
  )
 
 
 
 
 
 
 
 
 
1
  import unittest
2
+ from unittest.mock import patch
3
  import numpy as np
4
  from pysr import pysr, get_hof, best, best_tex, best_callable, best_row
5
+ from pysr.sr import run_feature_selection, _handle_feature_selection, _yesno
6
  import sympy
7
  from sympy import lambdify
8
  import pandas as pd
 
148
  np.testing.assert_array_equal(
149
  np.sort(selected_X, axis=1), np.sort(X[:, [2, 3]], axis=1)
150
  )
151
+
152
+ class TestHelperFunctions(unittest.TestCase):
153
+
154
+ @patch('builtins.input', side_effect=['y', 'n'])
155
+ def test_yesno(self, mock_input):
156
+ # Assert that the yes/no function correctly deals with y/n
157
+ self.assertEqual(_yesno("Test"), True)
158
+ self.assertEqual(_yesno("Test"), False)