import unittest
import numpy as np

from robo.fmin import random_search, bayesian_optimization


def objective(x):
    y = (x - 0.5) ** 2
    return y


class TestFminInterface(unittest.TestCase):

    def setUp(self):
        self.lower = np.zeros([1])
        self.upper = np.ones([1])

    def test_random_search(self):
        res = random_search(objective_function=objective, lower=self.lower, upper=self.upper)
        assert len(res["x_opt"].shape) == 1
        assert res["x_opt"].shape[0] == 1
        assert res["x_opt"] >= 0
        assert res["x_opt"] <= 1

    def test_bayesian_optimization(self):
        res = bayesian_optimization(objective_function=objective,
                                    lower=self.lower,
                                    upper=self.upper,
                                    num_iterations=10)
        assert len(res["x_opt"].shape) == 1
        assert res["x_opt"].shape[0] == 1
        assert res["x_opt"] >= 0
        assert res["x_opt"] <= 1

if __name__ == "__main__":
    unittest.main()

@@ -143,16 +143,19 @@
 return y
+%5B0%5D
 %0A%0A%0Aclass
@@ -376,16 +376,44 @@
 jective,
+%0A                           
  lower=s
@@ -439,17 +439,17 @@
 lf.upper
-)
+,
 %0A       
@@ -453,43 +453,45 @@
     
-assert len(res%5B%22x_opt%22%5D.shape) == 1
+                    num_iterations=3)
 %0A   
@@ -494,32 +494,36 @@
 %0A        assert 
+len(
 res%5B%22x_opt%22%5D.sha
@@ -514,33 +514,25 @@
 res%5B%22x_opt%22%5D
-.shape%5B0%5D
+)
  == 1%0A      
@@ -532,32 +532,41 @@
 %0A        assert 
+np.array(
 res%5B%22x_opt%22%5D %3E= 
@@ -553,32 +553,33 @@
 ray(res%5B%22x_opt%22%5D
+)
  %3E= 0%0A        as
@@ -575,32 +575,41 @@
 %0A        assert 
+np.array(
 res%5B%22x_opt%22%5D %3C= 
@@ -600,24 +600,25 @@
 res%5B%22x_opt%22%5D
+)
  %3C= 1%0A%0A    d
@@ -820,32 +820,78 @@
 per=self.upper,%0A
+                                    n_init=2,%0A
                 
@@ -929,10 +929,9 @@
 ons=
-10
+3
 )%0A  
@@ -959,22 +959,16 @@
 %22x_opt%22%5D
-.shape
 ) == 1%0A 
@@ -985,50 +985,17 @@
 ert 
-res%5B%22x_opt%22%5D.shape%5B0%5D == 1%0A        assert 
+np.array(
 res%5B
@@ -1002,16 +1002,17 @@
 %22x_opt%22%5D
+)
  %3E= 0%0A  
@@ -1024,16 +1024,25 @@
  assert 
+np.array(
 res%5B%22x_o
@@ -1045,16 +1045,17 @@
 %22x_opt%22%5D
+)
  %3C= 1%0A%0Ai

fix unittest