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)