Skip to content

Commit

Permalink
Merge pull request #70 from t-bittarn/master
Browse files Browse the repository at this point in the history
Add must_not for bool search query
  • Loading branch information
vrcmarcos authored Jun 18, 2021
2 parents efa6b1b + 4c9292a commit 4ab4b24
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 0 deletions.
25 changes: 25 additions & 0 deletions elasticmock/fake_elasticsearch.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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')

Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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:
Expand Down
12 changes: 12 additions & 0 deletions tests/fake_elasticsearch/test_search.py
Original file line number Diff line number Diff line change
Expand Up @@ -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})
Expand Down

0 comments on commit 4ab4b24

Please sign in to comment.