Add additional tests for ON CONFLICT DO NOTHING duplicate rows filtering

Author: Photonios

psqlextra/query.py

@@ -41,6 +41,14 @@
     QuerySetBase = QuerySet
 
 
+def peek_iterator(iterable):
+    try:
+        first = next(iterable)
+    except StopIteration:
+        return None
+    return list(chain([first], iterable))
+
+
 class PostgresQuerySet(QuerySetBase, Generic[TModel]):
     """Adds support for PostgreSQL specifics."""
 
@@ -177,14 +185,7 @@ def bulk_insert(
         if rows is None:
             return []
 
-        def peek(iterable):
-            try:
-                first = next(iterable)
-            except StopIteration:
-                return None
-            return list(chain([first], iterable))
-
-        rows = peek(iter(rows))
+        rows = peek_iterator(iter(rows))
 
         if not rows:
             return []

tests/test_on_conflict_nothing.py

@@ -170,24 +170,26 @@ def test_on_conflict_nothing_foreign_key_by_id():
     assert obj1.data == "some data"
 
 
-def test_on_conflict_nothing_duplicate_rows():
+@pytest.mark.parametrize(
+    "rows,expected_row_count",
+    [
+        ([dict(amount=1), dict(amount=1)], 1),
+        (iter([dict(amount=1), dict(amount=1)]), 1),
+        ((row for row in [dict(amount=1), dict(amount=1)]), 1),
+        ([], 0),
+        (iter([]), 0),
+        ((row for row in []), 0),
+    ],
+)
+def test_on_conflict_nothing_duplicate_rows(rows, expected_row_count):
     """Tests whether duplicate rows are filtered out when doing a insert
     NOTHING and no error is raised when the list of rows contains
     duplicates."""
 
     model = get_fake_model({"amount": models.IntegerField(unique=True)})
 
-    rows = [dict(amount=1), dict(amount=1)]
-
-    inserted_rows = model.objects.on_conflict(
-        ["amount"], ConflictAction.NOTHING
-    ).bulk_insert(rows)
-
-    assert len(inserted_rows) == 1
-
-    rows = iter([dict(amount=2), dict(amount=2)])
     inserted_rows = model.objects.on_conflict(
         ["amount"], ConflictAction.NOTHING
     ).bulk_insert(rows)
 
-    assert len(inserted_rows) == 1
+    assert len(inserted_rows) == expected_row_count