Merge pull request #1 from TheoResources/master

count method

elasticmock/fake_elasticsearch.py

@@ -109,6 +109,37 @@ def get_source(self, index, doc_type, id, params=None):
         document = self.get(index=index, doc_type=doc_type, id=id, params=params)
         return document.get('_source')
 
+    @query_params('_source', '_source_exclude', '_source_include',
+                  'allow_no_indices', 'analyze_wildcard', 'analyzer', 'default_operator',
+                  'df', 'expand_wildcards', 'explain', 'fielddata_fields', 'fields',
+                  'from_', 'ignore_unavailable', 'lenient', 'lowercase_expanded_terms',
+                  'preference', 'q', 'request_cache', 'routing', 'scroll', 'search_type',
+                  'size', 'sort', 'stats', 'suggest_field', 'suggest_mode',
+                  '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()
+
+        i = 0
+        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
+                i += 1
+        result = {
+            'count': i,
+            '_shards': {
+                'successful': 1,
+                'failed': 0,
+                'total': 1
+            }
+        }
+
+        return result
+
     @query_params('_source', '_source_exclude', '_source_include',
                   'allow_no_indices', 'analyze_wildcard', 'analyzer', 'default_operator',
                   'df', 'expand_wildcards', 'explain', 'fielddata_fields', 'fields',

tests/test_elasticmock.py

@@ -10,7 +10,6 @@
 
 
 class TestFakeElasticsearch(TestCase):
-
     @elasticmock
     def setUp(self):
         self.es = elasticsearch.Elasticsearch(hosts=[{'host': 'localhost', 'port': 9200}])
@@ -94,6 +93,14 @@ def test_should_raise_notfounderror_when_search_for_unexistent_index(self):
         with self.assertRaises(NotFoundError):
             self.es.search(index=self.index_name)
 
+    def test_should_return_count_for_indexed_documents_on_index(self):
+        index_quantity = 0
+        for i in range(0, index_quantity):
+            self.es.index(index='index_{0}'.format(i), doc_type=self.doc_type, body={'data': 'test_{0}'.format(i)})
+
+        count = self.es.count()
+        self.assertEqual(index_quantity, count.get('count'))
+
     def test_should_return_all_documents(self):
         index_quantity = 10
         for i in range(0, index_quantity):