updates to generator and also add test cases

Author: jtsodapop

labelbox/data/generator.py

@@ -70,7 +70,7 @@ def fill_queue(self):
                 self.queue.put(value)
         except Exception as e:
             self.queue.put(
-                ValueError("Unexpected exception while filling queue. %r", e))
+                ValueError(f"Unexpected exception while filling queue: {e}"))
         finally:
             self.queue.put(None)
 
@@ -92,8 +92,8 @@ def __next__(self) -> Any:
                         thread.join()
                     raise StopIteration
                 value = self.queue.get()
-        if isinstance(value, Exception):
-            raise value
+            if isinstance(value, Exception):
+                raise value
         else:
             value = self._process(next(self._data))
         return value

tests/data/test_prefetch_generator.py

@@ -0,0 +1,34 @@
+import pytest
+from labelbox.data.generator import PrefetchGenerator
+from random import random
+
+
+class ChildClassGenerator(PrefetchGenerator):
+
+    def __init__(self, examples, num_executors=1):
+        super().__init__(data=examples, num_executors=num_executors)
+
+    def _process(self, value):
+        num = random()
+        if num < .2:
+            raise ValueError("Randomized value error")
+        return value
+
+
+amount = (i for i in range(50))
+
+
+def test_single_thread_generator():
+    generator = ChildClassGenerator(amount, num_executors=1)
+
+    with pytest.raises(ValueError):
+        for _ in range(51):
+            next(generator)
+
+
+def test_multi_thread_generator():
+    generator = ChildClassGenerator(amount, num_executors=4)
+
+    with pytest.raises(ValueError):
+        for _ in range(51):
+            next(generator)