// Path: src/main/java/org/elasticsearch/ml/training/ModelTrainer.java
// interface TrainingSession {
//     AggregationBuilder trainingRequest();
//     String model(SearchResponse response);
// }


import org.elasticsearch.ResourceNotFoundException;
import org.elasticsearch.Version;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.cluster.ClusterName;
import org.elasticsearch.cluster.ClusterState;
import org.elasticsearch.cluster.metadata.AliasMetaData;
import org.elasticsearch.cluster.metadata.IndexMetaData;
import org.elasticsearch.cluster.metadata.MappingMetaData;
import org.elasticsearch.cluster.metadata.MetaData;
import org.elasticsearch.cluster.service.ClusterService;
import org.elasticsearch.common.UUIDs;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.index.IndexNotFoundException;
import org.elasticsearch.index.query.MatchAllQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.TermQueryBuilder;
import org.elasticsearch.indices.query.IndicesQueriesRegistry;
import org.elasticsearch.ml.training.ModelTrainer.TrainingSession;
import org.elasticsearch.plugins.SearchPlugin.QuerySpec;
import org.elasticsearch.search.SearchRequestParsers;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.test.ESTestCase;
import org.hamcrest.Matcher;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.instanceOf;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
                .putMapping(mappingMetaData)
                .putAlias(AliasMetaData.builder("just_me_alias"))
                .putAlias(AliasMetaData.builder("other_and_me_alias"))
                .settings(MINIMAL_INDEX_SETTINGS);
        IndexMetaData.Builder otherIndexMetaData = IndexMetaData.builder("other_index")
                .putMapping(mappingMetaData)
                .putAlias(AliasMetaData.builder("other_and_me_alias"))
                .settings(MINIMAL_INDEX_SETTINGS);
        MetaData.Builder metaData = MetaData.builder().put(indexMetaData).put(otherIndexMetaData);
        when(clusterService.state()).thenReturn(ClusterState.builder(new ClusterName("test")).metaData(metaData).build());

        AggregationBuilder aggregationBuilder = mock(AggregationBuilder.class);
        SearchRequestBuilder searchRequestBuilder = mockSearchRequestBuilder.apply(aggregationBuilder);

        Client client = mock(Client.class);

        when(client.prepareSearch()).thenReturn(searchRequestBuilder);
        when(client.prepareSearch()).thenReturn(searchRequestBuilder);

        SearchResponse searchResponse = mock(SearchResponse.class);
        doAnswer(invocationOnMock -> {
            @SuppressWarnings("unchecked") ActionListener<SearchResponse> listener =
                    (ActionListener<SearchResponse>) invocationOnMock.getArguments()[0];
            listener.onResponse(searchResponse);
            return null;
        }).when(searchRequestBuilder).execute(any());

        ModelTrainer modelTrainer = mock(ModelTrainer.class);
        when(modelTrainer.modelType()).thenReturn("mock");
        ModelTrainers modelTrainers = new ModelTrainers(Collections.singletonList(modelTrainer));
// Path: src/main/java/org/elasticsearch/ml/training/ModelTrainer.java
// interface TrainingSession {
//     AggregationBuilder trainingRequest();
//     String model(SearchResponse response);
// }


// Path: src/test/java/org/elasticsearch/ml/training/TrainingServiceTests.java
import org.elasticsearch.ResourceNotFoundException;
import org.elasticsearch.Version;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.cluster.ClusterName;
import org.elasticsearch.cluster.ClusterState;
import org.elasticsearch.cluster.metadata.AliasMetaData;
import org.elasticsearch.cluster.metadata.IndexMetaData;
import org.elasticsearch.cluster.metadata.MappingMetaData;
import org.elasticsearch.cluster.metadata.MetaData;
import org.elasticsearch.cluster.service.ClusterService;
import org.elasticsearch.common.UUIDs;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.index.IndexNotFoundException;
import org.elasticsearch.index.query.MatchAllQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.TermQueryBuilder;
import org.elasticsearch.indices.query.IndicesQueriesRegistry;
import org.elasticsearch.ml.training.ModelTrainer.TrainingSession;
import org.elasticsearch.plugins.SearchPlugin.QuerySpec;
import org.elasticsearch.search.SearchRequestParsers;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.test.ESTestCase;
import org.hamcrest.Matcher;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.instanceOf;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

                .putMapping(mappingMetaData)
                .putAlias(AliasMetaData.builder("just_me_alias"))
                .putAlias(AliasMetaData.builder("other_and_me_alias"))
                .settings(MINIMAL_INDEX_SETTINGS);
        IndexMetaData.Builder otherIndexMetaData = IndexMetaData.builder("other_index")
                .putMapping(mappingMetaData)
                .putAlias(AliasMetaData.builder("other_and_me_alias"))
                .settings(MINIMAL_INDEX_SETTINGS);
        MetaData.Builder metaData = MetaData.builder().put(indexMetaData).put(otherIndexMetaData);
        when(clusterService.state()).thenReturn(ClusterState.builder(new ClusterName("test")).metaData(metaData).build());

        AggregationBuilder aggregationBuilder = mock(AggregationBuilder.class);
        SearchRequestBuilder searchRequestBuilder = mockSearchRequestBuilder.apply(aggregationBuilder);

        Client client = mock(Client.class);

        when(client.prepareSearch()).thenReturn(searchRequestBuilder);
        when(client.prepareSearch()).thenReturn(searchRequestBuilder);

        SearchResponse searchResponse = mock(SearchResponse.class);
        doAnswer(invocationOnMock -> {
            @SuppressWarnings("unchecked") ActionListener<SearchResponse> listener =
                    (ActionListener<SearchResponse>) invocationOnMock.getArguments()[0];
            listener.onResponse(searchResponse);
            return null;
        }).when(searchRequestBuilder).execute(any());

        ModelTrainer modelTrainer = mock(ModelTrainer.class);
        when(modelTrainer.modelType()).thenReturn("mock");
        ModelTrainers modelTrainers = new ModelTrainers(Collections.singletonList(modelTrainer));
        TrainingSession trainingSession = mock(TrainingSession.class);