Python iterators example: A random number iterator that doesn't repeat elements.

random_iterator.py

from random import randint


class RandomNumberIterator(object):
    def __init__(self, smallest=0, largest=1000):
        self.served = set()
        self.smallest = smallest
        self.largest = largest

    def __iter__(self):
        return self

    def __next__(self):
        if len(self.served) == (self.largest - self.smallest + 1):
            raise StopIteration()

        while True:
            number = randint(self.smallest, self.largest)
            if number not in self.served:
                self.served.add(number)
                return number

    next = __next__

tests.py

import unittest

from random_iterator import RandomNumberIterator

class RandomIteratorTestCase(unittest.TestCase):

    def test_iterator_finishes(self):
        it = RandomNumberIterator(0, 1)
        next(it)
        next(it)

        with self.assertRaises(StopIteration):
            next(it)

    def test_iterator_returns_correct_elements(self):
        it = RandomNumberIterator(0, 1)

        elems = set([])
        elems.add(next(it))
        elems.add(next(it))

        self.assertEqual(elems, set([0, 1]))