From edee5a8de6afb34a9247de2bb73ab66397a6e573 Mon Sep 17 00:00:00 2001 From: François Freitag Date: Thu, 1 Jun 2017 16:56:51 -0400 Subject: Fixed #27639 -- Added chunk_size parameter to QuerySet.iterator(). --- tests/queries/test_iterator.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 tests/queries/test_iterator.py (limited to 'tests') diff --git a/tests/queries/test_iterator.py b/tests/queries/test_iterator.py new file mode 100644 index 0000000000..56f42c2191 --- /dev/null +++ b/tests/queries/test_iterator.py @@ -0,0 +1,39 @@ +import datetime +from unittest import mock + +from django.db.models.sql.compiler import cursor_iter +from django.test import TestCase + +from .models import Article + + +class QuerySetIteratorTests(TestCase): + itersize_index_in_mock_args = 3 + + @classmethod + def setUpTestData(cls): + Article.objects.create(name='Article 1', created=datetime.datetime.now()) + Article.objects.create(name='Article 2', created=datetime.datetime.now()) + + def test_iterator_invalid_chunk_size(self): + for size in (0, -1): + with self.subTest(size=size): + with self.assertRaisesMessage(ValueError, 'Chunk size must be strictly positive.'): + Article.objects.iterator(chunk_size=size) + + def test_default_iterator_chunk_size(self): + qs = Article.objects.iterator() + with mock.patch('django.db.models.sql.compiler.cursor_iter', side_effect=cursor_iter) as cursor_iter_mock: + next(qs) + self.assertEqual(cursor_iter_mock.call_count, 1) + mock_args, _mock_kwargs = cursor_iter_mock.call_args + self.assertEqual(mock_args[self.itersize_index_in_mock_args], 2000) + + def test_iterator_chunk_size(self): + batch_size = 3 + qs = Article.objects.iterator(chunk_size=batch_size) + with mock.patch('django.db.models.sql.compiler.cursor_iter', side_effect=cursor_iter) as cursor_iter_mock: + next(qs) + self.assertEqual(cursor_iter_mock.call_count, 1) + mock_args, _mock_kwargs = cursor_iter_mock.call_args + self.assertEqual(mock_args[self.itersize_index_in_mock_args], batch_size) -- cgit v1.3