path: root/core/src/test/java/org/elasticsearch/search/rescore/QueryRescoreBuilderTests.java

/*
 * Licensed to Elasticsearch under one or more contributor
 * license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright
 * ownership. Elasticsearch licenses this file to you under
 * the Apache License, Version 2.0 (the "License"); you may
 * not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package org.elasticsearch.search.rescore;

import org.apache.lucene.search.Query;
import org.elasticsearch.ElasticsearchParseException;
import org.elasticsearch.Version;
import org.elasticsearch.cluster.metadata.IndexMetaData;
import org.elasticsearch.common.ParsingException;
import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.xcontent.NamedXContentRegistry;
import org.elasticsearch.common.xcontent.ToXContent;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentFactory;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.common.xcontent.json.JsonXContent;
import org.elasticsearch.index.IndexSettings;
import org.elasticsearch.index.mapper.ContentPath;
import org.elasticsearch.index.mapper.MappedFieldType;
import org.elasticsearch.index.mapper.Mapper;
import org.elasticsearch.index.mapper.TextFieldMapper;
import org.elasticsearch.index.query.MatchAllQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryShardContext;
import org.elasticsearch.search.SearchModule;
import org.elasticsearch.search.rescore.QueryRescorer.QueryRescoreContext;
import org.elasticsearch.test.ESTestCase;
import org.elasticsearch.test.IndexSettingsModule;
import org.junit.AfterClass;
import org.junit.BeforeClass;

import java.io.IOException;

import static java.util.Collections.emptyList;
import static org.elasticsearch.test.EqualsHashCodeTestUtils.checkEqualsAndHashCode;

public class QueryRescoreBuilderTests extends ESTestCase {

    private static final int NUMBER_OF_TESTBUILDERS = 20;
    private static NamedWriteableRegistry namedWriteableRegistry;
    private static NamedXContentRegistry xContentRegistry;

    /**
     * setup for the whole base test class
     */
    @BeforeClass
    public static void init() {
        SearchModule searchModule = new SearchModule(Settings.EMPTY, false, emptyList());
        namedWriteableRegistry = new NamedWriteableRegistry(searchModule.getNamedWriteables());
        xContentRegistry = new NamedXContentRegistry(searchModule.getNamedXContents());
    }

    @AfterClass
    public static void afterClass() throws Exception {
        namedWriteableRegistry = null;
        xContentRegistry = null;
    }

    /**
     * Test serialization and deserialization of the rescore builder
     */
    public void testSerialization() throws IOException {
        for (int runs = 0; runs < NUMBER_OF_TESTBUILDERS; runs++) {
            RescoreBuilder<?> original = randomRescoreBuilder();
            RescoreBuilder<?> deserialized = copy(original);
            assertEquals(deserialized, original);
            assertEquals(deserialized.hashCode(), original.hashCode());
            assertNotSame(deserialized, original);
        }
    }

    /**
     * Test equality and hashCode properties
     */
    public void testEqualsAndHashcode() throws IOException {
        for (int runs = 0; runs < NUMBER_OF_TESTBUILDERS; runs++) {
            checkEqualsAndHashCode(randomRescoreBuilder(), this::copy, QueryRescoreBuilderTests::mutate);
        }
    }

    private RescoreBuilder<?> copy(RescoreBuilder<?> original) throws IOException {
        return copyWriteable(original, namedWriteableRegistry,
                namedWriteableRegistry.getReader(RescoreBuilder.class, original.getWriteableName()));
    }

    /**
     *  creates random rescorer, renders it to xContent and back to new instance that should be equal to original
     */
    public void testFromXContent() throws IOException {
        for (int runs = 0; runs < NUMBER_OF_TESTBUILDERS; runs++) {
            RescoreBuilder<?> rescoreBuilder = randomRescoreBuilder();
            XContentBuilder builder = XContentFactory.contentBuilder(randomFrom(XContentType.values()));
            if (randomBoolean()) {
                builder.prettyPrint();
            }
            rescoreBuilder.toXContent(builder, ToXContent.EMPTY_PARAMS);
            XContentBuilder shuffled = shuffleXContent(builder);


            XContentParser parser = createParser(shuffled);
            parser.nextToken();
            RescoreBuilder<?> secondRescoreBuilder = RescoreBuilder.parseFromXContent(parser);
            assertNotSame(rescoreBuilder, secondRescoreBuilder);
            assertEquals(rescoreBuilder, secondRescoreBuilder);
            assertEquals(rescoreBuilder.hashCode(), secondRescoreBuilder.hashCode());
        }
    }

    /**
     * test that build() outputs a {@link RescoreSearchContext} that has the same properties
     * than the test builder
     */
    public void testBuildRescoreSearchContext() throws ElasticsearchParseException, IOException {
        final long nowInMillis = randomNonNegativeLong();
        Settings indexSettings = Settings.builder()
                .put(IndexMetaData.SETTING_VERSION_CREATED, Version.CURRENT).build();
        IndexSettings idxSettings = IndexSettingsModule.newIndexSettings(randomAlphaOfLengthBetween(1, 10), indexSettings);
        // shard context will only need indicesQueriesRegistry for building Query objects nested in query rescorer
        QueryShardContext mockShardContext = new QueryShardContext(0, idxSettings, null, null, null, null, null, xContentRegistry(),
                null, null, () -> nowInMillis) {
            @Override
            public MappedFieldType fieldMapper(String name) {
                TextFieldMapper.Builder builder = new TextFieldMapper.Builder(name);
                return builder.build(new Mapper.BuilderContext(idxSettings.getSettings(), new ContentPath(1))).fieldType();
            }
        };

        for (int runs = 0; runs < NUMBER_OF_TESTBUILDERS; runs++) {
            QueryRescorerBuilder rescoreBuilder = randomRescoreBuilder();
            QueryRescoreContext rescoreContext = rescoreBuilder.build(mockShardContext);
            int expectedWindowSize = rescoreBuilder.windowSize() == null ? QueryRescoreContext.DEFAULT_WINDOW_SIZE :
                rescoreBuilder.windowSize().intValue();
            assertEquals(expectedWindowSize, rescoreContext.window());
            Query expectedQuery = QueryBuilder.rewriteQuery(rescoreBuilder.getRescoreQuery(), mockShardContext).toQuery(mockShardContext);
            assertEquals(expectedQuery, rescoreContext.query());
            assertEquals(rescoreBuilder.getQueryWeight(), rescoreContext.queryWeight(), Float.MIN_VALUE);
            assertEquals(rescoreBuilder.getRescoreQueryWeight(), rescoreContext.rescoreQueryWeight(), Float.MIN_VALUE);
            assertEquals(rescoreBuilder.getScoreMode(), rescoreContext.scoreMode());
        }
    }

    public void testRescoreQueryNull() throws IOException {
        IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> new QueryRescorerBuilder((QueryBuilder) null));
        assertEquals("rescore_query cannot be null", e.getMessage());
    }

    /**
     * test parsing exceptions for incorrect rescorer syntax
     */
    public void testUnknownFieldsExpection() throws IOException {

        String rescoreElement = "{\n" +
                "    \"window_size\" : 20,\n" +
                "    \"bad_rescorer_name\" : { }\n" +
                "}\n";
        XContentParser parser = createParser(rescoreElement);
        try {
            RescoreBuilder.parseFromXContent(parser);
            fail("expected a parsing exception");
        } catch (ParsingException e) {
            assertEquals("rescore doesn't support rescorer with name [bad_rescorer_name]", e.getMessage());
        }

        rescoreElement = "{\n" +
                "    \"bad_fieldName\" : 20\n" +
                "}\n";
        parser = createParser(rescoreElement);
        try {
            RescoreBuilder.parseFromXContent(parser);
            fail("expected a parsing exception");
        } catch (ParsingException e) {
            assertEquals("rescore doesn't support [bad_fieldName]", e.getMessage());
        }

        rescoreElement = "{\n" +
                "    \"window_size\" : 20,\n" +
                "    \"query\" : [ ]\n" +
                "}\n";
        parser = createParser(rescoreElement);
        try {
            RescoreBuilder.parseFromXContent(parser);
            fail("expected a parsing exception");
        } catch (ParsingException e) {
            assertEquals("unexpected token [START_ARRAY] after [query]", e.getMessage());
        }

        rescoreElement = "{ }";
        parser = createParser(rescoreElement);
        try {
            RescoreBuilder.parseFromXContent(parser);
            fail("expected a parsing exception");
        } catch (ParsingException e) {
            assertEquals("missing rescore type", e.getMessage());
        }

        rescoreElement = "{\n" +
                "    \"window_size\" : 20,\n" +
                "    \"query\" : { \"bad_fieldname\" : 1.0  } \n" +
                "}\n";
        parser = createParser(rescoreElement);
        try {
            RescoreBuilder.parseFromXContent(parser);
            fail("expected a parsing exception");
        } catch (IllegalArgumentException e) {
            assertEquals("[query] unknown field [bad_fieldname], parser not found", e.getMessage());
        }

        rescoreElement = "{\n" +
                "    \"window_size\" : 20,\n" +
                "    \"query\" : { \"rescore_query\" : { \"unknown_queryname\" : { } } } \n" +
                "}\n";
        parser = createParser(rescoreElement);
        try {
            RescoreBuilder.parseFromXContent(parser);
            fail("expected a parsing exception");
        } catch (ParsingException e) {
            assertEquals("[query] failed to parse field [rescore_query]", e.getMessage());
        }

        rescoreElement = "{\n" +
                "    \"window_size\" : 20,\n" +
                "    \"query\" : { \"rescore_query\" : { \"match_all\" : { } } } \n"
                + "}\n";
        parser = createParser(rescoreElement);
        RescoreBuilder.parseFromXContent(parser);
    }

    /**
     * create a new parser from the rescorer string representation and reset context with it
     */
    private XContentParser createParser(String rescoreElement) throws IOException {
        XContentParser parser = createParser(JsonXContent.jsonXContent, rescoreElement);
        // move to first token, this is where the internal fromXContent
        assertTrue(parser.nextToken() == XContentParser.Token.START_OBJECT);
        return parser;
    }

    @Override
    protected NamedXContentRegistry xContentRegistry() {
        return xContentRegistry;
    }

    private static RescoreBuilder<?> mutate(RescoreBuilder<?> original) throws IOException {
        RescoreBuilder<?> mutation = ESTestCase.copyWriteable(original, namedWriteableRegistry, QueryRescorerBuilder::new);
        if (randomBoolean()) {
            Integer windowSize = original.windowSize();
            if (windowSize != null) {
                mutation.windowSize(windowSize + 1);
            } else {
                mutation.windowSize(randomIntBetween(0, 100));
            }
        } else {
            QueryRescorerBuilder queryRescorer = (QueryRescorerBuilder) mutation;
            switch (randomIntBetween(0, 3)) {
            case 0:
                queryRescorer.setQueryWeight(queryRescorer.getQueryWeight() + 0.1f);
                break;
            case 1:
                queryRescorer.setRescoreQueryWeight(queryRescorer.getRescoreQueryWeight() + 0.1f);
                break;
            case 2:
                QueryRescoreMode other;
                do {
                    other = randomFrom(QueryRescoreMode.values());
                } while (other == queryRescorer.getScoreMode());
                queryRescorer.setScoreMode(other);
                break;
            case 3:
                // only increase the boost to make it a slightly different query
                queryRescorer.getRescoreQuery().boost(queryRescorer.getRescoreQuery().boost() + 0.1f);
                break;
            default:
                throw new IllegalStateException("unexpected random mutation in test");
            }
        }
        return mutation;
    }

    /**
     * create random shape that is put under test
     */
    public static QueryRescorerBuilder randomRescoreBuilder() {
        QueryBuilder queryBuilder = new MatchAllQueryBuilder().boost(randomFloat())
                .queryName(randomAlphaOfLength(20));
        org.elasticsearch.search.rescore.QueryRescorerBuilder rescorer = new
                org.elasticsearch.search.rescore.QueryRescorerBuilder(queryBuilder);
        if (randomBoolean()) {
            rescorer.setQueryWeight(randomFloat());
        }
        if (randomBoolean()) {
            rescorer.setRescoreQueryWeight(randomFloat());
        }
        if (randomBoolean()) {
            rescorer.setScoreMode(randomFrom(QueryRescoreMode.values()));
        }
        if (randomBoolean()) {
            rescorer.windowSize(randomIntBetween(0, 100));
        }
        return rescorer;
    }
}