Unit Testing with OLS

test_ols.py

import numpy as np

def ols(x,y):
    x_mean = x.mean()
    y_mean = y.mean()
    beta_1 = np.multiply((x - x_mean), (y - y_mean)).sum() / np.square((x - x_mean)).sum()
    beta_0 = y_mean - (beta_1 * x_mean)
    
    def lin_model(new_x, slope=beta_1, intercept=beta_0):
        return new_x * slope + intercept

    return lin_model

def test_positive_slope():
	x = np.array([0, 1])
	y = np.array([0, 1])
	# slope to be 1 with intercept 0
	ln_model = ols(x, y)

	assert ln_model(2) == 2

def test_negative_slope():
	x = np.array([0, 1])
	y = np.array([0, -1])
	# slope to be -1 with intercept 0
	ln_model = ols(x, y)

	assert ln_model(2) == -2

def test_positive_intercept():
	x = np.array([0, 1])
	y = np.array([2, 4])
	# slope to be 2 with intercept 2
	ln_model = ols(x, y)

	assert ln_model(3) == 8

def test_negative_intercept():
	x = np.array([0, 1])
	y = np.array([-2, -4])
	# slope to be -2 with intercept -2
	ln_model = ols(x, y)

	assert ln_model(3) == -8

def test_vertical_line():
	x = np.array([1, 1])
	y = np.array([2, 3])
	# slope to be infinite
	ln_model = ols(x, y)

	assert ln_model(3) != ln_model(3)

def test_horizontal_line():
	x = np.array([0, 1])
	y = np.array([3, 3])
	# slope to be 0 with intercept 3
	ln_model = ols(x, y)

	assert ln_model(3) ==  3