Merge pull request #70 from t-bittarn/master

Add must_not for bool search query
vrcmarcos · Jun 18, 2021 · 4ab4b24 · 4ab4b24
2 parents efa6b1b + 4c9292a
commit 4ab4b24
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 0 deletions.
diff --git 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
@@ -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})