Search and count accept list of index names as 'index' params

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

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()