From 9f5ff07675c05ab32f7cc35b54ad305cdba6d082 Mon Sep 17 00:00:00 2001
From: Benjamin Dauvergne <bdauvergne@entrouvert.com>
Date: Thu, 9 Jul 2020 11:58:51 +0200
Subject: [PATCH] wcs_api: use WcsApi.batch_size for FormDatas batch size
 (#44970)

---
 tests/test_wcs.py   | 19 ++++++++++++++++++-
 wcs_olap/wcs_api.py | 12 ++++++------
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/tests/test_wcs.py b/tests/test_wcs.py
index e1c5a6a..70de2fe 100644
--- a/tests/test_wcs.py
+++ b/tests/test_wcs.py
@@ -19,6 +19,8 @@ import json
 import pytest
 import pathlib
 
+import unittest.mock as mock
+
 import requests
 import httmock
 
@@ -35,7 +37,22 @@ def test_wcs_fixture(wcs, postgres_db, tmpdir, olap_cmd, caplog):
             c.execute('CREATE TABLE olap_temp.coin()')
             c.execute('COMMIT')
 
-    olap_cmd()
+    class MockSession(requests.Session):
+        mocks = []
+
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self.get = mock.MagicMock(wraps=self.get)
+            self.mocks.append(self.get)
+
+    with mock.patch('requests.Session', MockSession):
+        olap_cmd()
+    call_args_list = MockSession.mocks[-1].call_args_list
+    url_with_limits = [call_args[0][0] for call_args in MockSession.mocks[-1].call_args_list
+                       if 'limit' in call_args[0][0]]
+    assert url_with_limits, call_args_list
+    for url_with_limit in url_with_limits:
+        assert 'limit=500&' in url_with_limit
 
     expected_schema = [
         ('agent', 'id'),
diff --git a/wcs_olap/wcs_api.py b/wcs_olap/wcs_api.py
index b175842..3862faa 100644
--- a/wcs_olap/wcs_api.py
+++ b/wcs_olap/wcs_api.py
@@ -218,12 +218,12 @@ class Schema(BaseObject):
 
 
 class FormDatas(object):
-    def __init__(self, wcs_api, formdef, full=False, anonymize=False, batch=1000):
+    def __init__(self, wcs_api, formdef, full=False, anonymize=False):
         self.wcs_api = wcs_api
         self.formdef = formdef
         self._full = full
         self.anonymize = anonymize
-        self.batch = batch
+        self.batch_size = wcs_api.batch_size
 
     def __getitem__(self, slice_or_id):
         # get batch of forms
@@ -272,19 +272,19 @@ class FormDatas(object):
 
     def batched(self, batch):
         forms = copy.copy(self)
-        forms.batch = batch
+        forms.batch_size = batch
         return forms
 
     def __iter__(self):
         start = 0
         while True:
             empty = True
-            for formdef in self[start:start + self.batch]:
+            for formdef in self[start:start + self.batch_size]:
                 empty = False
                 yield formdef
             if empty:
                 break
-            start += self.batch
+            start += self.batch_size
 
     def __len__(self):
         return len(list((o for o in self)))
@@ -533,7 +533,6 @@ class WcsApi(object):
     def __init__(self, url, email=None, name_id=None, batch_size=1000,
                  session=None, logger=None, orig=None, key=None, verify=True):
         self.url = url
-        self.batch_size = batch_size
         self.email = email
         self.name_id = name_id
         self.requests = session or requests.Session()
@@ -541,6 +540,7 @@ class WcsApi(object):
         self.orig = orig
         self.key = key
         self.verify = verify
+        self.batch_size = batch_size
 
     def _build_url(self, url_parts):
         url = self.url
-- 
2.26.2