Method test_partitions

test/ext/asyncio/test_engine.py:1323–1391

(
        self, async_engine, filter_, yield_per, yield_per_type
    )

    @testing.combinations("method", "opt", argnames="yield_per_type")
    @async_test
    async def test_partitions(
        self, async_engine, filter_, yield_per, yield_per_type
    ):
        users = self.tables.users
        async with async_engine.connect() as conn:
            stmt = select(users)
            if yield_per and yield_per_type == "opt":
                stmt = stmt.execution_options(yield_per=yield_per)
            result = await conn.stream(stmt)

            if filter_ == "mappings":
                result = result.mappings()
            elif filter_ == "scalars":
                result = result.scalars(1)

            if yield_per and yield_per_type == "method":
                result = result.yield_per(yield_per)

            check_result = []

            # stream() sets stream_results unconditionally
            assert isinstance(
                result._real_result.cursor_strategy,
                _cursor.BufferedRowCursorFetchStrategy,
            )

            if yield_per:
                partition_size = yield_per

                eq_(result._real_result.cursor_strategy._bufsize, yield_per)

                async for partition in result.partitions():
                    check_result.append(partition)
            else:
                eq_(result._real_result.cursor_strategy._bufsize, 5)

                partition_size = 5
                async for partition in result.partitions(partition_size):
                    check_result.append(partition)

            ranges = [
                (i, min(20, i + partition_size))
                for i in range(1, 21, partition_size)
            ]

            if filter_ == "mappings":
                eq_(
                    check_result,
                    [
                        [
                            {"user_id": i, "user_name": "name%d" % i}
                            for i in range(a, b)
                        ]
                        for (a, b) in ranges
                    ],
                )
            elif filter_ == "scalars":
                eq_(