diff --git a/elasticmock/fake_elasticsearch.py b/elasticmock/fake_elasticsearch.py index 5afa5e2..ee09176 100644 --- a/elasticmock/fake_elasticsearch.py +++ b/elasticmock/fake_elasticsearch.py @@ -118,10 +118,7 @@ def get_source(self, index, doc_type, id, params=None): 'suggest_size', 'suggest_text', 'terminate_after', 'timeout', 'track_scores', 'version') def count(self, index=None, doc_type=None, body=None, params=None): - if index is not None and index not in self.__documents_dict: - raise NotFoundError(404, 'IndexMissingException[[{0}] missing]'.format(index)) - - searchable_indexes = [index] if index is not None else self.__documents_dict.keys() + searchable_indexes = self._normalize_index_to_list(index) i = 0 for searchable_index in searchable_indexes: @@ -149,10 +146,7 @@ def count(self, index=None, doc_type=None, body=None, params=None): 'suggest_size', 'suggest_text', 'terminate_after', 'timeout', 'track_scores', 'version') def search(self, index=None, doc_type=None, body=None, params=None): - if index is not None and index not in self.__documents_dict: - raise NotFoundError(404, 'IndexMissingException[[{0}] missing]'.format(index)) - - searchable_indexes = [index] if index is not None else self.__documents_dict.keys() + searchable_indexes = self._normalize_index_to_list(index) matches = [] for searchable_index in searchable_indexes: @@ -235,3 +229,22 @@ def suggest(self, body, index=None, params=None): } ] return result_dict + + def _normalize_index_to_list(self, index): + # Ensure to have a list of index + if index is None: + searchable_indexes = self.__documents_dict.keys() + elif isinstance(index, str) or isinstance(index, unicode): + searchable_indexes = [index] + elif isinstance(index, list): + searchable_indexes = index + else: + # Is it the correct exception to use ? + raise ValueError("Invalid param 'index'") + + # Check index(es) exists + for searchable_index in searchable_indexes: + if searchable_index not in self.__documents_dict: + raise NotFoundError(404, 'IndexMissingException[[{0}] missing]'.format(searchable_index)) + + return searchable_indexes diff --git a/tests/test_elasticmock.py b/tests/test_elasticmock.py index 180617d..1821867 100644 --- a/tests/test_elasticmock.py +++ b/tests/test_elasticmock.py @@ -212,6 +212,22 @@ def test_should_return_suggestions(self): ], }, suggestion) + def test_should_search_in_multiple_indexes(self): + self.es.index(index='groups', doc_type='groups', body={'budget': 1000}) + self.es.index(index='users', doc_type='users', body={'name': 'toto'}) + self.es.index(index='pcs', doc_type='pcs', body={'model': 'macbook'}) + + result = self.es.search(index=['users', 'pcs']) + self.assertEqual(2, result.get('hits').get('total')) + + def test_should_count_in_multiple_indexes(self): + self.es.index(index='groups', doc_type='groups', body={'budget': 1000}) + self.es.index(index='users', doc_type='users', body={'name': 'toto'}) + self.es.index(index='pcs', doc_type='pcs', body={'model': 'macbook'}) + + result = self.es.count(index=['users', 'pcs']) + self.assertEqual(2, result.get('count')) + if __name__ == '__main__': unittest.main()