def test_generator_from_wikipedia(self):
generator = GeneratorFromWikipedia()
i = 0
while i < 100:
img, lbl = next(generator)
self.assertTrue(img.size[1] == 32, "Shape is not right")
i += 1
def test_generator_from_wikipedia_stops(self):
generator = GeneratorFromWikipedia(count=1)
next(generator)
self.assertRaises(StopIteration, generator.next)