def test_get_arm_initial_stage(self):
alg = algorithms.UCB1(3, {})
alg.update(0, 1)
alg.update(2, 1)
# Pull Arm 1 which has not been pulled.
self.assertEqual(1, alg.get_arm(1))
def setUp(self):
super(UCB1Test, self).setUp()
self._alg = algorithms.UCB1(2, {})
self._alg.set_state(np.array([10, 10]), np.array([10, 0]))