diff --git a/tests/unit/calibration/models/test_polynomial_model.py b/tests/unit/calibration/models/test_polynomial_model.py index c771f19..9fc7245 100644 --- a/tests/unit/calibration/models/test_polynomial_model.py +++ b/tests/unit/calibration/models/test_polynomial_model.py @@ -1,69 +1,59 @@ -import unittest -from functools import cached_property -from unittest.mock import patch, ANY +from unittest.mock import ANY import numpy as np +import pytest from depiction.calibration.models.polynomial_model import PolynomialModel -class TestPolynomialModel(unittest.TestCase): - def setUp(self) -> None: - self.mock_coef = [1, 2, 3] +@pytest.fixture() +def mock_coef(): + return [1, 2, 3] - @cached_property - def mock_model(self): - return PolynomialModel(coef=self.mock_coef) - def test_coef(self) -> None: - np.testing.assert_array_equal( - np.array([1, 2, 3]), - self.mock_model.coef, - ) +@pytest.fixture() +def mock_model(mock_coef): + return PolynomialModel(coef=mock_coef) - def test_is_zero_when_false(self) -> None: - self.assertFalse(self.mock_model.is_zero) - def test_is_zero_when_true(self) -> None: - self.assertTrue(PolynomialModel.zero().is_zero) - self.mock_coef = [0.0, 0, 0] - self.assertTrue(self.mock_model.is_zero) +def test_coef(mock_model) -> None: + np.testing.assert_array_equal(mock_model.coef, np.array([1, 2, 3])) - def test_degree(self) -> None: - self.assertEqual(2, self.mock_model.degree) - def test_predict(self) -> None: - np.testing.assert_array_equal( - np.array([6.0, 11]), - self.mock_model.predict([1.0, 2]), - ) +def test_is_zero_when_false(mock_model) -> None: + assert not mock_model.is_zero - def test_identity(self) -> None: - np.testing.assert_array_equal( - np.array([0, 1]), - PolynomialModel.identity().coef, - ) - def test_zero(self) -> None: - np.testing.assert_array_equal( - np.array([0, 0]), - PolynomialModel.zero().coef, - ) +def test_is_zero_when_true() -> None: + assert PolynomialModel.zero().is_zero + assert PolynomialModel([0.0, 0, 0]).is_zero - @patch("numpy.polyfit") - def test_fit_lsq(self, mock_polyfit) -> None: - self.mock_model_type = "poly_5" - mock_x = np.array([100, 200, 300]) - mock_y = np.array([1, 2, 3]) - mock_polyfit.return_value = np.array([5, 7]) - model = PolynomialModel.fit_lsq(mock_x, mock_y, degree=5) +def test_degree(mock_model) -> None: + assert mock_model.degree == 2 - np.testing.assert_array_equal(np.array([5, 7]), model.coef) - mock_polyfit.assert_called_once_with(ANY, ANY, deg=5) - np.testing.assert_array_equal(mock_x, mock_polyfit.call_args[0][0]) - np.testing.assert_array_equal(mock_y, mock_polyfit.call_args[0][1]) +def test_predict(mock_model) -> None: + np.testing.assert_array_equal(mock_model.predict([1.0, 2]), np.array([6.0, 11])) -if __name__ == "__main__": - unittest.main() + +def test_identity() -> None: + np.testing.assert_array_equal(np.array([0, 1]), PolynomialModel.identity().coef) + + +def test_zero() -> None: + np.testing.assert_array_equal(np.array([0, 0]), PolynomialModel.zero().coef) + + +def test_fit_lsq(mocker) -> None: + mock_polyfit = mocker.patch("numpy.polyfit") + mock_x = np.array([100, 200, 300]) + mock_y = np.array([1, 2, 3]) + mock_polyfit.return_value = np.array([5, 7]) + + model = PolynomialModel.fit_lsq(mock_x, mock_y, degree=5) + + np.testing.assert_array_equal(np.array([5, 7]), model.coef) + mock_polyfit.assert_called_once_with(ANY, ANY, deg=5) + np.testing.assert_array_equal(mock_x, mock_polyfit.call_args[0][0]) + np.testing.assert_array_equal(mock_y, mock_polyfit.call_args[0][1])