def setUp(self):
self.res = UniformReservoir(15)
class TestUniformReservoir(TestCase):
def setUp(self):
self.res = UniformReservoir(15)
def test_add_15_elements(self):
for i in range(15):
self.res.update(i)
self.assertEqual(len(self.res), 15)
def test_add_30_elements(self):
for i in range(30):
self.res.update(i)
self.assertEqual(len(self.res), 30)
self.assertEqual(len(self.res._res), 15)
@patch('caliper.reservoir.randint')
def test_randint_called_with_count_minus_one(self, randint):
for i in range(15):
self.res.update(i)
randint.return_value = 1
self.res.update(42)
randint.assert_called_once_with(0, 14)
@patch('caliper.reservoir.randint')
def test_full_reservoir_insert_in_correct_position(self, randint):
for _ in range(15):
self.res.update(0)
randint.return_value = 5
self.res.update(42)
randint.assert_called_once_with(0, 14)
self.assertEqual(self.res._res[5], 42)
randint.reset_mock()
randint.return_value = 0
self.res.update(21)
randint.assert_called_once_with(0, 15)
self.assertEqual(self.res._res[0], 21)
randint.reset_mock()
randint.return_value = 14
self.res.update(1337)
randint.assert_called_once_with(0, 16)
self.assertEqual(self.res._res[-1], 1337)
randint.reset_mock()
def test_full_reservoir_ignores_index_too_large(self):
for _ in range(30):
self.res.update(0)
with patch('caliper.reservoir.randint') as randint:
randint.return_value = 20
self.res.update(42)
randint.assert_called_once_with(0, 29)
self.assertFalse(42 in self.res._res)