diff --git a/elasticmock/fake_elasticsearch.py b/elasticmock/fake_elasticsearch.py index 3286b6d..921195a 100644 --- a/elasticmock/fake_elasticsearch.py +++ b/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', diff --git a/tests/test_elasticmock.py b/tests/test_elasticmock.py index 0950753..60d1939 100644 --- a/tests/test_elasticmock.py +++ b/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):