Improve search and count

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,37 +146,35 @@ 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:
             for document in self.__documents_dict[searchable_index]:
                 if doc_type is not None and document.get('_type') != doc_type:
                     continue
                 matches.append(document)
+
         result = {
             'hits': {
                 'total': len(matches),
                 'max_score': 1.0
             },
             '_shards': {
-                'successful': 1,
+                # Simulate indexes with 1 shard each
+                'successful': len(searchable_indexes),
                 'failed': 0,
-                'total': 1
+                'total': len(searchable_indexes)
             },
             'took': 1,
             'timed_out': False
         }
 
-        if matches:
-            hits = []
-            for match in matches:
-                match['_score'] = 1.0
-                hits.append(match)
-            result['hits']['hits'] = hits
+        hits = []
+        for match in matches:
+            match['_score'] = 1.0
+            hits.append(match)
+        result['hits']['hits'] = hits
 
         return result
 
@@ -234,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

@@ -101,6 +101,11 @@ def test_should_return_count_for_indexed_documents_on_index(self):
         count = self.es.count()
         self.assertEqual(index_quantity, count.get('count'))
 
+    def test_should_return_hits_hits_even_when_no_result(self):
+        search = self.es.search()
+        self.assertEqual(0, search.get('hits').get('total'))
+        self.assertListEqual([], search.get('hits').get('hits'))
+
     def test_should_return_all_documents(self):
         index_quantity = 10
         for i in range(0, index_quantity):
@@ -207,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()