import unittest
import numbers
import numpy as np
from problem import generate_objective_function
from robot_arm import RobotArm


class TestObjectiveFunction(unittest.TestCase):
    def setUp(self):
        self.thetas = np.array((0.5, 4, 2, 5, 3, 6,))
        self.thetas.setflags(write=False)
        lengths = (1, 1,)
        destinations = ((0, 0, 0), (0, 0, 0))
        self.robot_arm = RobotArm(lengths, destinations)
        self.objective = generate_objective_function(self.robot_arm)

    def test_theta_size(self):
        self.assertEqual(self.thetas.shape, (2*3,))

    def test_return_type(self):
        self.assertIsInstance(self.objective(self.thetas), numbers.Number)

    def test_correct_objective_value(self):
        self.assertEqual(self.objective(self.thetas), 31 / 4)

    def test_immutability(self):
        original = self.thetas
        self.objective(self.thetas)
        np.testing.assert_equal(original, self.thetas)

@@ -9,16 +9,42 @@
 nittest%0A
+from numpy import testing%0A
 import n
@@ -116,16 +116,54 @@
 function
+, generate_objective_gradient_function
 %0Afrom ro
@@ -914,16 +914,23 @@
 f.thetas
+.copy()
 %0A       
@@ -1001,20 +1001,651 @@
 ginal, self.thetas)%0A
+%0A    def test_objective_value2(self):%0A        robot_arm = RobotArm(%0A            lengths=(1, 2,),%0A            destinations=((1, -1,), (1, 1,),)%0A        )%0A        thetas = np.array((0, np.pi/2, np.pi/2, np.pi/4,))%0A%0A        objective_func = generate_objective_function(robot_arm)%0A        objective_value = objective_func(thetas)%0A        self.assertEqual(objective_value, np.pi**2 * 5/16)%0A%0A        objective_gradient_func = generate_objective_gradient_function(robot_arm)%0A        objective_gradient = objective_gradient_func(thetas)%0A        testing.assert_array_equal(objective_gradient, np.array((-np.pi, np.pi/2, np.pi, -np.pi/2,)))%0A

Add additinal objective tests