diff --git a/tests/integration/utils/test_sqlalchemy_exec_utils.py b/tests/integration/utils/test_sqlalchemy_exec_utils.py index bcc1c91a9..b549ad0f4 100644 --- a/tests/integration/utils/test_sqlalchemy_exec_utils.py +++ b/tests/integration/utils/test_sqlalchemy_exec_utils.py @@ -10,7 +10,8 @@ class SomeTable(Base): __tablename__ = "some_table" - pk = sqlalchemy.Column(sqlalchemy.Integer, primary_key=True) + pk = sqlalchemy.Column(sqlalchemy.Integer, primary_key=True, autoincrement=False) + key = sqlalchemy.Column(sqlalchemy.Integer) foo = sqlalchemy.Column(sqlalchemy.String) @@ -21,9 +22,9 @@ def test_fetch_table_in_batches_unique(engine): table_size = 15 batch_size = 6 - table_data = [(i, f"foo{i}") for i in range(table_size)] + table_data = [(i, i, f"foo{i}") for i in range(table_size)] - engine.setup([SomeTable(pk=row[0], foo=row[1]) for row in table_data]) + engine.setup([SomeTable(pk=row[0], key=row[1], foo=row[2]) for row in table_data]) table = SomeTable.__table__ @@ -31,10 +32,36 @@ def test_fetch_table_in_batches_unique(engine): results = fetch_table_in_batches( connection.execute, table, - table.c.pk, + table.c.key, key_is_unique=True, batch_size=batch_size, ) results = list(results) assert results == table_data + + +def test_fetch_table_in_batches_nonunique(engine): + if engine.name == "in_memory": + pytest.skip("SQL tests do not apply to in-memory engine") + + batch_size = 6 + repeats = [1, 2, 3, 4, 5, 0, 5, 4, 3, 2, 1] + keys = [key for key, n in enumerate(repeats) for _ in range(n)] + table_data = [(i, key, f"foo{i}") for i, key in enumerate(keys)] + + engine.setup([SomeTable(pk=row[0], key=row[1], foo=row[2]) for row in table_data]) + + table = SomeTable.__table__ + + with engine.sqlalchemy_engine().connect() as connection: + results = fetch_table_in_batches( + connection.execute, + table, + table.c.key, + key_is_unique=False, + batch_size=batch_size, + ) + results = sorted(results) + + assert results == table_data