From 37e8afcf8cf0f37bef49a3b21a65c7b877dd8fc5 Mon Sep 17 00:00:00 2001
From: Mike McLean <mikem@redhat.com>
Date: Fri, 3 May 2019 21:11:48 -0400
Subject: [PATCH] support batch option in MultiCallSession

---
 koji/__init__.py                         | 26 ++++++++++++++++--------
 tests/test_lib/test_multicall_session.py | 22 ++++++++++++++++++++
 2 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/koji/__init__.py b/koji/__init__.py
index ed53aa56..f2d4c67e 100644
--- a/koji/__init__.py
+++ b/koji/__init__.py
@@ -2928,9 +2928,10 @@ class MultiCallSession(object):
 
     """Manages a single multicall, acts like a session"""
 
-    def __init__(self, session, strict=False):
+    def __init__(self, session, strict=False, batch=None):
         self._session = session
-        self.strict = strict
+        self._strict = strict
+        self._batch = batch
         self._calls = []
 
     def __getattr__(self, name):
@@ -2949,7 +2950,7 @@ class MultiCallSession(object):
         """compatibility wrapper for _callMethod"""
         return self._callMethod(name, args, opts)
 
-    def call_all(self, strict=None):
+    def call_all(self, strict=None, batch=None):
         """Perform all calls in a single multicall
 
         Returns a the hub's multiCall result, which is a list of results for
@@ -2959,17 +2960,26 @@ class MultiCallSession(object):
         """
 
         if strict is None:
-            strict = self.strict
+            strict = self._strict
+        if batch is None:
+            batch = self._batch
 
         if len(self._calls) == 0:
             return []
 
         calls = self._calls
         self._calls = []
-        args = ([c.format() for c in calls],)
-        results = self._session._callMethod('multiCall', args, {})
-        for call, result in zip(calls, results):
-            call._result = result
+        if batch:
+            batches = [calls[i:i+batch] for i in range(0, len(calls), batch)]
+        else:
+            batches = [calls]
+        results = []
+        for calls in batches:
+            args = ([c.format() for c in calls],)
+            _results = self._session._callMethod('multiCall', args, {})
+            for call, result in zip(calls, _results):
+                call._result = result
+            results.extend(_results)
         if strict:
             # check for faults and raise first one
             for entry in results:
diff --git a/tests/test_lib/test_multicall_session.py b/tests/test_lib/test_multicall_session.py
index 5d15d580..15b9142d 100644
--- a/tests/test_lib/test_multicall_session.py
+++ b/tests/test_lib/test_multicall_session.py
@@ -31,3 +31,25 @@ class TestNewMultiCall(unittest.TestCase):
         for i in range(10):
             self.assertEqual(calls[i]['methodName'], "echo")
             self.assertEqual(calls[i]['params'], (i,))
+
+    def test_batch_multicall(self):
+        with self.session.multicall(batch=10) as m:
+            ret = {}
+            for i in range(42):
+                ret[i] = m.echo(i)
+
+        # should be 5 batches
+        self.assertEqual(self._callMethod.call_count, 5)
+        i = 0
+        for args, kwargs in self._callMethod.call_args_list:
+            self.assertEqual(kwargs, {})
+            self.assertEqual(args[0], 'multiCall')
+            self.assertEqual(args[2], {})
+            _calls = args[1]
+            if not isinstance(_calls, tuple) or len(_calls) != 1:
+                raise Exception('multiCall args not wrapped in singleton')
+            calls = _calls[0]
+            for call in calls:
+                self.assertEqual(call['methodName'], "echo")
+                self.assertEqual(call['params'], (i,))
+                i += 1