def test_random_stimuli_with_distance(self):
samples = [Sample([x]) for x in xrange(10)] * 10
chooser = RandomStimuliChooser(None, True, 3)
env = Environment(samples, chooser)
for _ in xrange(10):
sort = sorted([x.get_values()[0] for x in env.get_stimuli(4)])
self.assertEqual([0, 3, 6, 9], sort)
chooser = RandomStimuliChooser(None, True, 5)
env = Environment(samples, chooser)
for _ in xrange(10):
sort = sorted([x.get_values()[0] for x in env.get_stimuli(2)])
self.assertEqual(len(sort), 2)
self.assertTrue(sort[0] < 5)
self.assertTrue(sort[1] >= 5)
self.assertRaises(Exception, chooser.get_stimuli, samples, 100)
