Fix test for timed_window_unique node

Author: bdewilde

streamz/tests/test_core.py

@@ -333,58 +333,51 @@ def read_from_q():
     assert end - start >= 0.2
 
 
-# TODO: looks like gen_test and pytest.mark.parametrize don't play nicely together:
-# TypeError: test_timed_window_unique() missing 5 required positional arguments: 'interval', 'key', 'keep', 'elements', and 'exp_result'
-# I think gen_test() needs to be modified for this to work, but unsure how
-# TODO: fix tests' exp_result values and the logic inside, once it runs without ^ error
-# @gen_test()
-# @pytest.mark.parametrize(
-#     "interval,key,keep,elements,exp_result",
-#     [
-#         (0.2, sz.identity, "first", [1, 2, 1, 3, 1, 3, 3, 2], [(1, 2, 3), (1, 3, 2)]),
-#         (0.2, sz.identity, "last", [1, 2, 1, 3, 1, 3, 3, 2], [(2, 1, 3), (1, 3, 2)]),
-#         (
-#             0.02,
-#             len,
-#             "last",
-#             ["f", "fo", "f", "foo", "f", "foo", "foo", "fo"],
-#             [("fo", "f", "foo"), ("f", "foo", "fo")],
-#         ),
-#         (
-#             0.02,
-#             "id",
-#             "first",
-#             [{"id": 0, "foo": "bar"}, {"id": 0, "foo": "baz"}, {"id": 1, "foo": "bat"}],
-#             [({"id": 0, "foo": "bar"}, {"id": 1, "foo": "bat"})],
-#         ),
-#         (
-#             0.02,
-#             "id",
-#             "last",
-#             [{"id": 0, "foo": "bar"}, {"id": 0, "foo": "baz"}, {"id": 1, "foo": "bat"}],
-#             [({"id": 0, "foo": "baz"}, {"id": 1, "foo": "bat"})],
-#         ),
-#     ]
-# )
-# def test_timed_window_unique(interval, key, keep, elements, exp_result):
-#     import pdb; pdb.set_trace()
-#     source = Stream(asynchronous=True)
-#     a = source.timed_window_unique(interval, key, keep)
-
-#     assert a.loop is IOLoop.current()
-#     L = a.sink_to_list()
-
-#     for ele in elements:
-#         yield source.emit(ele)
-#         yield gen.sleep(0.006)
-#     yield gen.sleep(a.interval)
-
-#     assert L
-#     assert all(len(x) <= 3 for x in L)
-#     assert any(len(x) >= 2 for x in L)
-
-#     yield gen.sleep(0.1)
-#     assert not L[-1]
+@gen_test()
+def test_timed_window_unique():
+    tests = [
+        (0.05, sz.identity, "first", [1, 2, 1, 3, 1, 3, 3, 2], [(1, 2, 3)]),
+        (0.05, sz.identity, "last", [1, 2, 1, 3, 1, 3, 3, 2], [(1, 3, 2)]),
+        (
+            0.05,
+            len,
+            "last",
+            ["f", "fo", "f", "foo", "f", "foo", "foo", "fo"],
+            [("f", "foo", "fo")],
+        ),
+        (
+            0.05,
+            "id",
+            "first",
+            [{"id": 0, "foo": "bar"}, {"id": 1, "foo": "bat"}, {"id": 0, "foo": "baz"}],
+            [({"id": 0, "foo": "bar"}, {"id": 1, "foo": "bat"})],
+        ),
+        (
+            0.05,
+            "id",
+            "last",
+            [{"id": 0, "foo": "bar"}, {"id": 1, "foo": "bat"}, {"id": 0, "foo": "baz"}],
+            [({"id": 1, "foo": "bat"}, {"id": 0, "foo": "baz"})],
+        ),
+    ]
+    for interval, key, keep, elements, exp_result in tests:
+        source = Stream(asynchronous=True)
+        a = source.timed_window_unique(interval, key, keep)
+
+        assert a.loop is IOLoop.current()
+        L = a.sink_to_list()
+
+        for ele in elements:
+            yield source.emit(ele)
+        yield gen.sleep(a.interval)
+
+        assert L
+        assert all(wi in elements for window in L for wi in window)
+        assert sum(1 for window in L for _ in window) <= len(elements)
+        assert L == exp_result
+
+        yield gen.sleep(a.interval)
+        assert not L[-1]
 
 
 @gen_test()