From 0c192b5dafebaeb9063b20cda04f558b62752014 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carlos=20G=C3=A1lvez=20Mart=C3=ADnez?= Date: Wed, 20 Jan 2021 17:14:14 +0100 Subject: [PATCH 1/2] Add multi_match Allow evaluation for query_type MUST --- elasticmock/fake_elasticsearch.py | 29 +++++++++++++++++++++++++ tests/fake_elasticsearch/test_search.py | 26 ++++++++++++++++++++++ 2 files changed, 55 insertions(+) diff --git a/elasticmock/fake_elasticsearch.py b/elasticmock/fake_elasticsearch.py index 63986db..5926677 100644 --- a/elasticmock/fake_elasticsearch.py +++ b/elasticmock/fake_elasticsearch.py @@ -31,6 +31,7 @@ class QueryType: RANGE = 'RANGE' SHOULD = 'SHOULD' MINIMUM_SHOULD_MATCH = 'MINIMUM_SHOULD_MATCH' + MULTI_MATCH = 'MULTI_MATCH' @staticmethod def get_query_type(type_str): @@ -54,6 +55,8 @@ def get_query_type(type_str): return QueryType.SHOULD elif type_str == 'minimum_should_match': return QueryType.MINIMUM_SHOULD_MATCH + elif type_str == 'multi_match': + return QueryType.MULTI_MATCH else: raise NotImplementedError(f'type {type_str} is not implemented for QueryType') @@ -95,6 +98,10 @@ def _evaluate_for_query_type(self, document): return self._evaluate_for_compound_query_type(document) elif self.type == QueryType.FILTER: return self._evaluate_for_compound_query_type(document) + elif self.type == QueryType.MUST: + return self._evaluate_for_compound_query_type(document) + elif self.type == QueryType.MULTI_MATCH: + return self._evaluate_for_multi_match_query_type(document) else: raise NotImplementedError('Fake query evaluation not implemented for query type: %s' % self.type) @@ -125,6 +132,25 @@ def _evaluate_for_field(self, document, ignore_case): break return return_val + def _evaluate_for_fields(self, document): + doc_source = document['_source'] + return_val = False + value = self.condition.get('query') + if not value: + return return_val + fields = self.condition.get('fields', []) + for field in fields: + return_val = self._compare_value_for_field( + doc_source, + field, + value, + True + ) + if return_val: + break + + return return_val + def _evaluate_for_range_query_type(self, document): for field, comparisons in self.condition.items(): doc_val = document['_source'] @@ -180,6 +206,9 @@ def _evaluate_for_compound_query_type(self, document): return return_val + def _evaluate_for_multi_match_query_type(self, document): + return self._evaluate_for_fields(document) + def _compare_value_for_field(self, doc_source, field, value, ignore_case): if ignore_case and isinstance(value, str): value = value.lower() diff --git a/tests/fake_elasticsearch/test_search.py b/tests/fake_elasticsearch/test_search.py index 465f7fc..98aac2b 100644 --- a/tests/fake_elasticsearch/test_search.py +++ b/tests/fake_elasticsearch/test_search.py @@ -163,6 +163,32 @@ def test_query_on_nested_data(self): doc = response['hits']['hits'][0]['_source'] self.assertEqual(i, doc['id']) + + def test_search_with_bool_query_and_multi_match(self): + for i in range(0, 10): + self.es.index(index='index_for_search', doc_type=DOC_TYPE, body={ + 'data': 'test_{0}'.format(i) if i % 2 == 0 else None, + 'data2': 'test_{0}'.format(i) if (i+1) % 2 == 0 else None + }) + + search_body = { + "query": { + "bool": { + "must": { + "multi_match": { + "query": "test", + "fields": ["data", "data2"] + } + } + } + } + } + response = self.es.search(index='index_for_search', doc_type=DOC_TYPE, + body=search_body) + self.assertEqual(response['hits']['total'], 10) + hits = response['hits']['hits'] + self.assertEqual(len(hits), 10) + @parameterized.expand( [ ( From fb4f6ad51717fe1e9c6252cd9249ef01fa853650 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carlos=20G=C3=A1lvez=20Mart=C3=ADnez?= Date: Thu, 21 Jan 2021 10:54:34 +0100 Subject: [PATCH 2/2] Add msearch --- elasticmock/fake_elasticsearch.py | 32 ++++++++++++++++++++ tests/fake_elasticsearch/test_search.py | 39 +++++++++++++++++++++++++ 2 files changed, 71 insertions(+) diff --git a/elasticmock/fake_elasticsearch.py b/elasticmock/fake_elasticsearch.py index 5926677..63cc9c0 100644 --- a/elasticmock/fake_elasticsearch.py +++ b/elasticmock/fake_elasticsearch.py @@ -438,6 +438,38 @@ def count(self, index=None, doc_type=None, body=None, params=None, headers=None) def _get_fake_query_condition(self, query_type_str, condition): return FakeQueryCondition(QueryType.get_query_type(query_type_str), condition) + @query_params( + "ccs_minimize_roundtrips", + "max_concurrent_searches", + "max_concurrent_shard_requests", + "pre_filter_shard_size", + "rest_total_hits_as_int", + "search_type", + "typed_keys", + ) + def msearch(self, body, index=None, doc_type=None, params=None, headers=None): + def grouped(iterable): + if len(iterable) % 2 != 0: + raise Exception('Malformed body') + iterator = iter(iterable) + while True: + try: + yield (next(iterator)['index'], next(iterator)) + except StopIteration: + break + + responses = [] + took = 0 + for ind, query in grouped(body): + response = self.search(index=ind, body=query) + took += response['took'] + responses.append(response) + result = { + 'took': took, + 'responses': responses + } + 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/fake_elasticsearch/test_search.py b/tests/fake_elasticsearch/test_search.py index 98aac2b..348f82e 100644 --- a/tests/fake_elasticsearch/test_search.py +++ b/tests/fake_elasticsearch/test_search.py @@ -189,6 +189,45 @@ def test_search_with_bool_query_and_multi_match(self): hits = response['hits']['hits'] self.assertEqual(len(hits), 10) + def test_msearch(self): + for i in range(0, 10): + self.es.index(index='index_for_search1', doc_type=DOC_TYPE, body={ + 'data': 'test_{0}'.format(i) if i % 2 == 0 else None, + 'data2': 'test_{0}'.format(i) if (i+1) % 2 == 0 else None + }) + for i in range(0, 10): + self.es.index(index='index_for_search2', doc_type=DOC_TYPE, body={ + 'data': 'test_{0}'.format(i) if i % 2 == 0 else None, + 'data2': 'test_{0}'.format(i) if (i+1) % 2 == 0 else None + }) + + search_body = { + "query": { + "bool": { + "must": { + "multi_match": { + "query": "test", + "fields": ["data", "data2"] + } + } + } + } + } + body = [] + body.append({'index': 'index_for_search1'}) + body.append(search_body) + body.append({'index': 'index_for_search2'}) + body.append(search_body) + + result = self.es.msearch(index='index_for_search', body=body) + response1, response2 = result['responses'] + self.assertEqual(response1['hits']['total'], 10) + hits1 = response1['hits']['hits'] + self.assertEqual(len(hits1), 10) + self.assertEqual(response2['hits']['total'], 10) + hits2 = response2['hits']['hits'] + self.assertEqual(len(hits2), 10) + @parameterized.expand( [ (