lr-gd-test.py

class TestGradientDescent(unittest.TestCase):

    def test_correct_prediction(self):
      global X
      global y
      if len(X.shape) != 2:        
        X = X.reshape(X.shape[0], 1)
      w = fit(X, y)
      y_hat = predict(X, w).round()
      self.assertTrue((y_hat == y).all())