def testCountSimple(self): items = iterator.Counted(iter(["foo", "bar", "baz"])) self.assertEqual(items.count, 0) self.assertEqual(next(items), "foo") self.assertEqual(items.count, 1) self.assertEqual(next(items), "bar") self.assertEqual(items.count, 2) self.assertEqual(next(items), "baz") self.assertEqual(items.count, 3) with self.assertRaises(StopIteration): next(items)
def testResetSimple(self): items = iterator.Counted(iter(["foo", "bar", "baz"])) next(items) self.assertEqual(items.count, 1) items.Reset() self.assertEqual(items.count, 0) next(items) next(items) self.assertEqual(items.count, 2) items.Reset() self.assertEqual(items.count, 0)
def testCountInterrupted(self): def Generator(): yield "foo" raise ValueError() yield "bar" # pylint: disable=unreachable items = iterator.Counted(Generator()) self.assertEqual(items.count, 0) self.assertEqual(next(items), "foo") self.assertEqual(items.count, 1) with self.assertRaises(ValueError): next(items) # The counter should not increase if the item was not retrieved. self.assertEqual(items.count, 1)
def Run(self, args: rdf_timeline.TimelineArgs) -> None: """Executes the client action.""" entries = iterator.Counted(Walk(args.root)) for entry_batch in rdf_timeline.TimelineEntry.SerializeStream(entries): entry_batch_blob = rdf_protodict.DataBlob(data=entry_batch) self.SendReply(entry_batch_blob, session_id=self._TRANSFER_STORE_ID) entry_batch_blob_id = hashlib.sha256(entry_batch).digest() result = rdf_timeline.TimelineResult() result.entry_batch_blob_ids.append(entry_batch_blob_id) result.entry_count = entries.count self.SendReply(result) # Each result should contain information only about the number of entries # in the current batch, so after the results are sent we simply reset the # counter. entries.Reset()
def testCountEmpty(self): empty = iterator.Counted(iter([])) self.assertEqual(empty.count, 0) with self.assertRaises(StopIteration): next(empty)
def testResetEmpty(self): empty = iterator.Counted(iter([])) empty.Reset() self.assertEqual(empty.count, 0)