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)