def test_dumped_example_visitor(self):
example_id_manager = example_id_visitor.ExampleIdManager(
self.data_source, 0)
visitor = example_id_visitor.ExampleIdVisitor(example_id_manager)
expected_index = 0
for (index, example) in visitor:
self.assertEqual(index, expected_index)
self.assertEqual('{}'.format(index).encode(), example.example_id)
self.assertEqual(150000000 + index, example.event_time)
self.assertEqual(index, example.index)
expected_index += 1
self.assertEqual(self.end_index, index)
try:
visitor.seek(1 << 30)
except StopIteration:
self.assertTrue(True)
self.assertTrue(visitor.finished())
else:
self.assertTrue(False)
index, example = visitor.seek(500)
self.assertFalse(visitor.finished())
expected_index = 500
self.assertEqual(index, expected_index)
self.assertEqual('{}'.format(index).encode(), example.example_id)
self.assertEqual(150000000 + index, example.event_time)
self.assertEqual(index, example.index)
expected_index = 501
for (index, example) in visitor:
self.assertEqual(index, expected_index)
self.assertEqual('{}'.format(index).encode(), example.example_id)
self.assertEqual(150000000 + index, example.event_time)
self.assertEqual(index, example.index)
expected_index += 1
self.assertEqual(self.end_index, index)
def test_example_id_dumper(self):
example_id_dumper1 = example_id_dumper.ExampleIdDumperManager(
self.kvstore, self.data_source, 0, self.example_id_dump_options)
self.assertEqual(example_id_dumper1.get_next_index(), 0)
self._dump_example_ids(example_id_dumper1, 0, 10, 1024)
example_id_manager = \
example_id_visitor.ExampleIdManager(self.kvstore, self.data_source, 0, True)
last_dumped_index = example_id_manager.get_last_dumped_index()
self.assertEqual(last_dumped_index, 10240 - 1)
example_id_dumper2 = example_id_dumper.ExampleIdDumperManager(
self.kvstore, self.data_source, 0, self.example_id_dump_options)
self.assertEqual(example_id_dumper2.get_next_index(), 10240)
self._dump_example_ids(example_id_dumper2, 10 * 1024, 10, 1024)
last_dumped_index = example_id_manager.get_last_dumped_index()
self.assertEqual(last_dumped_index, 2 * 10240 - 1)