From 4c9292a360d4e636dae56bf2348635aba9d06ca7 Mon Sep 17 00:00:00 2001 From: "t@bittarn.com" Date: Fri, 18 Jun 2021 16:53:17 +0800 Subject: [PATCH] Add must_not for bool search query --- elasticmock/fake_elasticsearch.py | 25 +++++++++++++++++++++++++ tests/fake_elasticsearch/test_search.py | 12 ++++++++++++ 2 files changed, 37 insertions(+) diff --git a/elasticmock/fake_elasticsearch.py b/elasticmock/fake_elasticsearch.py index 6b580ff..1372952 100644 --- a/elasticmock/fake_elasticsearch.py +++ b/elasticmock/fake_elasticsearch.py @@ -35,6 +35,7 @@ class QueryType: SHOULD = 'SHOULD' MINIMUM_SHOULD_MATCH = 'MINIMUM_SHOULD_MATCH' MULTI_MATCH = 'MULTI_MATCH' + MUST_NOT = 'MUST_NOT' @staticmethod def get_query_type(type_str): @@ -60,6 +61,8 @@ def get_query_type(type_str): return QueryType.MINIMUM_SHOULD_MATCH elif type_str == 'multi_match': return QueryType.MULTI_MATCH + elif type_str == 'must_not': + return QueryType.MUST_NOT else: raise NotImplementedError(f'type {type_str} is not implemented for QueryType') @@ -107,6 +110,8 @@ def _evaluate_for_query_type(self, document): return self._evaluate_for_should_query_type(document) elif self.type == QueryType.MULTI_MATCH: return self._evaluate_for_multi_match_query_type(document) + elif self.type == QueryType.MUST_NOT: + return self._evaluate_for_must_not_query_type(document) else: raise NotImplementedError('Fake query evaluation not implemented for query type: %s' % self.type) @@ -211,6 +216,26 @@ def _evaluate_for_compound_query_type(self, document): return return_val + def _evaluate_for_must_not_query_type(self, document): + if isinstance(self.condition, dict): + for query_type, sub_query in self.condition.items(): + return_val = FakeQueryCondition( + QueryType.get_query_type(query_type), + sub_query + ).evaluate(document) + if return_val: + return False + elif isinstance(self.condition, list): + for sub_condition in self.condition: + for sub_condition_key in sub_condition: + return_val = FakeQueryCondition( + QueryType.get_query_type(sub_condition_key), + sub_condition[sub_condition_key] + ).evaluate(document) + if return_val: + return False + return True + def _evaluate_for_should_query_type(self, document): return_val = False for sub_condition in self.condition: diff --git a/tests/fake_elasticsearch/test_search.py b/tests/fake_elasticsearch/test_search.py index 7b71070..f31c64b 100644 --- a/tests/fake_elasticsearch/test_search.py +++ b/tests/fake_elasticsearch/test_search.py @@ -141,6 +141,18 @@ def test_search_with_bool_query(self): hits = response['hits']['hits'] self.assertEqual(len(hits), 1) + def test_search_with_must_not_query(self): + for i in range(0, 10): + self.es.index(index='index_for_search', doc_type=DOC_TYPE, body={'id': i}) + response = self.es.search(index='index_for_search', doc_type=DOC_TYPE, + body={'query': {'bool': { + 'filter': [{'terms': {'id': [1, 2]}}], + 'must_not': [{'term': {'id': 1}}], + }}}) + self.assertEqual(response['hits']['total']['value'], 1) + doc = response['hits']['hits'][0]['_source'] + self.assertEqual(2, doc['id']) + def test_search_with_terms_query(self): for i in range(0, 10): self.es.index(index='index_for_search', doc_type=DOC_TYPE, body={'id': i})