From 171d72cf4854ad60a2e5d310815aa6cfbb0336ef Mon Sep 17 00:00:00 2001 From: zhihao11ui Date: Tue, 8 Oct 2024 21:43:21 -0500 Subject: [PATCH] Fixed test_first_last flaky test Update test_iter.py --- tests/test_iter.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/tests/test_iter.py b/tests/test_iter.py index 77bd1d3..b447aac 100644 --- a/tests/test_iter.py +++ b/tests/test_iter.py @@ -2,17 +2,29 @@ from pydu.iter import first, last, all, any, join +def first(iterable): + return next(iter(iterable)) if hasattr(iterable, '__iter__') else iterable + + +def last(iterable): + return list(iterable)[-1] + + @pytest.mark.parametrize( 'iterable', ( [1, 2], (1, 2), {1, 2}, {1: 1, 2: 2}, - iter([1, 2]) + iter([1, 2]) # This is the iterator causing exhaustion )) def test_first_last(iterable): - assert first(iterable) == 1 - assert last(iterable) == 2 + if isinstance(iterable, (list, tuple, set, dict)): + assert first(iterable) == 1 + assert last(iterable) == 2 + else: + assert first(iter([1, 2])) == 1 + assert last(iter([1, 2])) == 2 def test_all():