def setUp(self): self._es = Elasticsearch() self._es_index = 'vistest' self._es_type = 'test' self._scroll_time = '1m' self._migrator = EsMigrator(self._es, self._es_index) self._visibility = Visibility(formalVisibility='U')
class TestVisibility(TestCase): def setUp(self): self._es = Elasticsearch() self._es_index = 'vistest' self._es_type = 'test' self._scroll_time = '1m' self._migrator = EsMigrator(self._es, self._es_index) self._visibility = Visibility(formalVisibility='U') def testUpdate00(self): self._assertUpdateBehavior(0, 0) def testUpdate11(self): self._assertUpdateBehavior(1, 1) def testUpdate15(self): self._assertUpdateBehavior(1, 5) def testUpdate55(self): self._assertUpdateBehavior(5, 5) def _assertUpdateBehavior(self, numBatches, batchSize, visibility=Visibility(formalVisibility='U')): batchSizes = [batchSize] * numBatches + [0] self._migrator._es.search = MagicMock(return_value={'_scroll_id': 0}) self._migrator._es.scroll = Mock(side_effect=TestVisibility.MockScrollSideEffect(batchSizes)) self._migrator._es.update = Mock() ### self._migrator.addVisibility({}, self._es_type, scroll_time=self._scroll_time, visibility=self._visibility) ### self._migrator._es.search.assert_called_with( doc_type=self._es_type, index=self._es_index, scroll=self._scroll_time, search_type='scan' ) self.assertEqual( [call(scroll=self._scroll_time, scroll_id=i) for i in range(numBatches+1)], self._migrator._es.scroll.call_args_list ) self.assertEqual( [call( body={'doc': {'ezbake_visibility': self._b64vis(visibility)}}, doc_type=self._es_type, id=i, index=self._es_index ) for i in range(batchSize)] * numBatches, self._migrator._es.update.call_args_list ) def _b64vis(self, visibility): return serialize_to_base64(visibility) class MockScrollSideEffect(): def __init__(self, batchSizes): self._batchSizes = batchSizes def __call__(self, *args, **kw): scroll_id = (kw['scroll_id'] if kw.has_key('scroll_id') else 0) batchSizes = self._batchSizes[scroll_id] print 'called with %d. batchSizes: %d' % (scroll_id, batchSizes) return { '_scroll_id': scroll_id + 1, 'hits' : {'hits': [{'_id':x} for x in range(batchSizes)]} } class MockUpdateSideEffect(): def __init__(self): pass def __call__(self, *args, **kw): pass