import unittest
import numpy as np
from numpy.testing import assert_array_almost_equal

from big_o import complexities


class TestComplexities(unittest.TestCase):

    def test_compute(self):
        x = np.linspace(10, 100, 100)
        y = 3.0 * x + 2.0
        linear = complexities.Linear()
        linear.fit(x, y)
        assert_array_almost_equal(linear.compute(x), y, 10)

    def test_not_fitted(self):
        linear = complexities.Linear()
        self.assertRaises(complexities.NotFittedError, linear.compute, 100)

    def test_str_includes_units(self):
        x = np.linspace(10, 100, 100)
        y = 3.0 * x + 2.0
        linear = complexities.Linear()
        linear.fit(x, y)
        linear_str = str(linear)
        assert '(sec)' in linear_str

@@ -198,117 +198,674 @@
     
-x = np.linspace(10, 100, 100)%0A        y = 3.0 * x + 2.0%0A        linear = complexities.Linear()%0A        linear
+desired = %5B%0A            (lambda x: 2.+x*0., complexities.Constant),%0A            (lambda x: 5.*x+3., complexities.Linear),%0A            (lambda x: 8.1*x**2.+0.9, complexities.Quadratic),%0A            (lambda x: 1.0*x**3+11.0, complexities.Cubic),%0A            (lambda x: 5.2*x**2.5, complexities.Polynomial),%0A            (lambda x: 8.5*np.log(x)+99.0, complexities.Logarithmic),%0A            (lambda x: 1.7*x*np.log(x)+2.74, complexities.Linearithmic),%0A            (lambda x: 3.14**x, complexities.Exponential)%0A        %5D%0A%0A        x = np.linspace(10, 100, 100)%0A        for f, class_ in desired:%0A            y = f(x)%0A            complexity = class_()%0A            complexity
 .fit
@@ -863,32 +863,36 @@
 exity.fit(x, y)%0A
+    
         assert_a
@@ -909,22 +909,26 @@
 t_equal(
-linear
+complexity
 .compute
@@ -937,16 +937,83 @@
 ), y, 10
+, %22compute() failed to match expected values for class %25r%22 %25 class_
 )%0A%0A    d

Add compute test cases for all complexity classes