fix tests

avihad · Jul 3, 2016 · e17a219 · e17a219
1 parent c850b95
commit e17a219
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/tests/test_elasticsearch.py b/tests/test_elasticsearch.py
@@ -286,7 +286,7 @@ def test_scan(self):
         scroll_result = yield self.es.scan(SOME_INDEX, SOME_DOC_TYPE, query=some_search_query, scroll=scroll_size)
         self.es.search.assert_called_once_with(index=SOME_INDEX, doc_type=SOME_DOC_TYPE, body=some_search_query,
                                                scroll=scroll_size, search_type='scan')
-        self.assertEqual(Scroller(self.es, some_search_result).__dict__, scroll_result.__dict__)
+        self.assertEqual(Scroller(self.es, some_search_result, scroll_size).__dict__, scroll_result.__dict__)
 
     @inlineCallbacks
     def test_count(self):

diff --git a/tests/test_scroller.py b/tests/test_scroller.py
@@ -8,14 +8,15 @@
 SOME_VALUE_1 = "SOME_VALUE_1"
 SOME_VALUE_2 = "SOME_VALUE_2"
 SOME_ID = "SOME_ID"
+SOME_SCROLL = "2m"
 
 
 class TestScroller(TestCase):
     @inlineCallbacks
     def test_scroll_return_results(self):
         expected_results = [{SOME_VALUE_1: SOME_VALUE_2}]
         some_results = self.wrap_good_result(expected_results, SOME_ID)
-        scroller = Scroller(MagicMock(), some_results)
+        scroller = Scroller(MagicMock(), some_results, SOME_SCROLL)
         results = yield scroller.next()
         self.assertEqual(expected_results, results)
 
@@ -27,11 +28,11 @@ def test_scroll_next_scroller_iteration_call_es(self):
         expected_results = [{SOME_VALUE_2: SOME_VALUE_1}]
         es = MagicMock()
         es.scroll = MagicMock(return_value=self.wrap_good_result(expected_results, scroll_id))
-        scroller = Scroller(es, some_results1)
+        scroller = Scroller(es, some_results1, SOME_SCROLL)
         yield scroller.next()
         results = yield scroller.next()
         self.assertEqual(results, expected_results)
-        es.scroll.assert_called_once_with(scroll_id)
+        es.scroll.assert_called_once_with(scroll_id, scroll=SOME_SCROLL)
 
     @inlineCallbacks
     def test_scroll_iterator(self):
@@ -42,7 +43,7 @@ def test_scroll_iterator(self):
         expected_result_2 = [{SOME_VALUE_2: SOME_VALUE_1}]
         es = MagicMock()
         es.scroll = MagicMock(return_value=self.wrap_good_result(expected_result_2, None))
-        scroller = Scroller(es, es_results_1)
+        scroller = Scroller(es, es_results_1, SOME_SCROLL)
 
         results = []
         for defer_results in scroller:
@@ -54,7 +55,7 @@ def test_scroll_iterator(self):
     def test_scroll_end_of_scan_when_scroll_id_is_none(self):
         expected_results = [{SOME_VALUE_1: SOME_VALUE_2}]
         some_results = self.wrap_good_result(expected_results, None)
-        scroller = Scroller(MagicMock(), some_results)
+        scroller = Scroller(MagicMock(), some_results, SOME_SCROLL)
         scroller.next()
         self.assertRaises(StopIteration, scroller.next)