Refactored dataloader tests

promise/promise.py

@@ -428,11 +428,14 @@ def on_error(error):
                 on_error,
             )
 
+            peek = Context.peek_context()
+            if peek:
+                peek.drain_queue()
+
             if self._trace:
                 # If we wait, we drain the queue of the
                 # callbacks waiting on the context exit
                 # so we avoid a blocking state
-                Context.peek_context().drain_queue()
                 self._trace.drain_queue()
 
             self._is_waiting = True

tests/test_dataloader.py

@@ -43,43 +43,8 @@ def call_fn(keys):
     assert values == []
 
 
-def test_batches_multiple_requests():
-    identity_loader, load_calls = id_loader()
-
-    @Promise.safe
-    def safe():
-        promise1 = identity_loader.load(1)
-        promise2 = identity_loader.load(2)
-        return promise1, promise2
-
-    promise1, promise2 = safe()
-    value1, value2 = Promise.all([promise1, promise2]).get()
-    assert value1 == 1
-    assert value2 == 2
-
-    assert load_calls == [[1, 2]]
-
-
-def test_batches_multiple_requests_two():
-    identity_loader, load_calls = id_loader()
-
-    @Promise.safe
-    def safe():
-        promise1 = identity_loader.load(1)
-        promise2 = identity_loader.load(2)
-        return Promise.all([promise1, promise2])
-
-    p = safe()
-    value1, value2 = p.get()
-
-    assert value1 == 1
-    assert value2 == 2
-
-    assert load_calls == [[1, 2]]
-
-
 @Promise.safe
-def test_batches_multiple_requests_safe():
+def test_batches_multiple_requests():
     identity_loader, load_calls = id_loader()
 
     promise1 = identity_loader.load(1)

tests/test_dataloader_extra.py

@@ -0,0 +1,65 @@
+from promise import Promise
+from promise.dataloader import DataLoader
+
+
+def id_loader(**options):
+    load_calls = []
+
+    def fn(keys):
+        load_calls.append(keys)
+        return Promise.resolve(keys)
+
+    identity_loader = DataLoader(fn, **options)
+    return identity_loader, load_calls
+
+
+def test_batches_multiple_requests():
+    identity_loader, load_calls = id_loader()
+
+    @Promise.safe
+    def safe():
+        promise1 = identity_loader.load(1)
+        promise2 = identity_loader.load(2)
+        return promise1, promise2
+
+    promise1, promise2 = safe()
+    value1, value2 = Promise.all([promise1, promise2]).get()
+    assert value1 == 1
+    assert value2 == 2
+
+    assert load_calls == [[1, 2]]
+
+
+def test_batches_multiple_requests_two():
+    identity_loader, load_calls = id_loader()
+
+    @Promise.safe
+    def safe():
+        promise1 = identity_loader.load(1)
+        promise2 = identity_loader.load(2)
+        return Promise.all([promise1, promise2])
+
+    p = safe()
+    value1, value2 = p.get()
+
+    assert value1 == 1
+    assert value2 == 2
+
+    assert load_calls == [[1, 2]]
+
+
+@Promise.safe
+def test_batches_multiple_requests_safe():
+    identity_loader, load_calls = id_loader()
+
+    promise1 = identity_loader.load(1)
+    promise2 = identity_loader.load(2)
+
+    p = Promise.all([promise1, promise2])
+
+    value1, value2 = p.get()
+
+    assert value1 == 1
+    assert value2 == 2
+
+    assert load_calls == [[1, 2]]