code: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF 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.apache.skywalking.library.elasticsearch;
  
  import org.apache.skywalking.library.elasticsearch.client.TemplateClient;
  import org.apache.skywalking.library.elasticsearch.requests.IndexRequest;
  import org.apache.skywalking.library.elasticsearch.requests.UpdateRequest;
  import org.apache.skywalking.library.elasticsearch.requests.search.Query;
  import org.apache.skywalking.library.elasticsearch.requests.search.Search;
  import org.apache.skywalking.library.elasticsearch.requests.search.aggregation.Aggregation;
  import org.apache.skywalking.library.elasticsearch.response.Document;
  import org.apache.skywalking.library.elasticsearch.response.Documents;
  import org.apache.skywalking.library.elasticsearch.response.IndexTemplate;
  import org.apache.skywalking.library.elasticsearch.response.Mappings;
  import org.apache.skywalking.library.elasticsearch.response.search.SearchResponse;
  import org.awaitility.Duration;
  import org.junit.jupiter.api.Tag;
  import org.junit.jupiter.params.ParameterizedTest;
  import org.junit.jupiter.params.provider.MethodSource;
  import org.testcontainers.elasticsearch.ElasticsearchContainer;
  import org.testcontainers.shaded.com.google.common.collect.ImmutableMap;
  import org.testcontainers.utility.DockerImageName;
  
  import java.util.Arrays;
  import java.util.Collection;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  
  import static org.assertj.core.api.Assertions.assertThat;
  import static org.awaitility.Awaitility.await;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertNotNull;
  import static org.junit.jupiter.api.Assertions.assertNull;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  
  @Tag("slow")
  public class ElasticSearchIT {
      public static Collection<Object[]> es() {
          // noinspection resource
          return Arrays.asList(new Object[][] {
              {
                  "ElasticSearch 6.3.2",
                  new ElasticsearchContainer(
                      DockerImageName.parse("docker.elastic.co/elasticsearch/elasticsearch-oss")
                                     .withTag("6.3.2"))
              },
              {
                  "ElasticSearch 7.4.2",
                  new ElasticsearchContainer(
                      DockerImageName.parse("docker.elastic.co/elasticsearch/elasticsearch-oss")
                                     .withTag("7.4.2"))
              },
              {
                  "ElasticSearch 7.8.0",
                  new ElasticsearchContainer(
                      DockerImageName.parse("docker.elastic.co/elasticsearch/elasticsearch-oss")
                                     .withTag("7.8.0"))
              },
              {
                  "ElasticSearch 7.15.0",
                  new ElasticsearchContainer(
                      DockerImageName.parse("elastic/elasticsearch")
                                     .withTag("7.15.0")
                                     .asCompatibleSubstituteFor(
                                         "docker.elastic.co/elasticsearch/elasticsearch-oss"))
              },
              {
                  "ElasticSearch 8.1.0",
                  new ElasticsearchContainer(
                      DockerImageName.parse("docker.elastic.co/elasticsearch/elasticsearch")
                          .withTag("8.1.0"))
                                  .withEnv("xpack.security.enabled", "false")
              },
              {
                  "OpenSearch 1.0.0",
                  new ElasticsearchContainer(
                      DockerImageName.parse("opensearchproject/opensearch")
                                     .withTag("1.0.0")
                                     .asCompatibleSubstituteFor(
                                         "docker.elastic.co/elasticsearch/elasticsearch-oss"))
                      .withEnv("plugins.security.disabled", "true")
                      .withStartupTimeout(java.time.Duration.ofMinutes(5))
              },
              {
                  "OpenSearch 2.4.0",
                  new ElasticsearchContainer(
                      DockerImageName.parse("opensearchproject/opensearch")
                                     .withTag("2.4.0")
                                     .asCompatibleSubstituteFor(
                                         "docker.elastic.co/elasticsearch/elasticsearch-oss"))
                      .withEnv("plugins.security.disabled", "true")
                      .withStartupTimeout(java.time.Duration.ofMinutes(5))
              }
          });
      }
  
      @ParameterizedTest(name = "version: {0}")
      @MethodSource("es")
      public void testTemplate(final String ignored,
                               final ElasticsearchContainer server) {
          server.start();
  
          final ElasticSearch client =
                  ElasticSearch.builder()
                          .endpoints(server.getHttpHostAddress())
                          .build();
          client.connect();
  
          final String name = "test-template";
          final TemplateClient templateClient = client.templates();
  
          final ImmutableMap<String, Object> properties = ImmutableMap.of(
              "metric_table", ImmutableMap.of("type", "keyword"),
              "service_id", ImmutableMap.of("type", "keyword")
          );
  
          final Mappings.Source sourceConf = new Mappings.Source();
          sourceConf.getExcludes().add("test");
  
          final Mappings mappings = Mappings.builder()
                                            .type("_doc")
                                            .properties(properties)
                                            .source(sourceConf)
                                            .build();
  
          assertThat(templateClient.createOrUpdate(name, ImmutableMap.of(), mappings, 0))
              .isTrue();
  
          assertThat(templateClient.exists(name)).isTrue();
  
          assertThat(templateClient.get(name))
              .isPresent()
              .map(IndexTemplate::getMappings)
              .map(Mappings::getProperties)
              .hasValue(mappings.getProperties());
          assertThat(templateClient.get(name))
              .isPresent()
              .map(IndexTemplate::getMappings)
              .map(Mappings::getSource)
              .map(Mappings.Source::getExcludes)
              .hasValue(mappings.getSource().getExcludes());
  
          server.close();
      }
  
      @ParameterizedTest(name = "version: {0}")
      @MethodSource("es")
      public void testIndex(final String ignored,
                            final ElasticsearchContainer server) {
          server.start();
  
          final ElasticSearch client =
                  ElasticSearch.builder()
                          .endpoints(server.getHttpHostAddress())
                          .build();
          client.connect();
  
          final String index = "test-index";
          assertFalse(client.index().exists(index));
          assertFalse(client.index().get(index).isPresent());
          assertTrue(client.index().create(index, null, null));
          assertTrue(client.index().exists(index));
          assertNotNull(client.index().get(index));
          assertTrue(client.index().delete(index));
          assertFalse(client.index().get(index).isPresent());
  
          server.close();
      }
  
      @ParameterizedTest(name = "version: {0}")
      @MethodSource("es")
      public void testDoc(final String ignored,
                          final ElasticsearchContainer server) {
          server.start();
  
          final ElasticSearch client =
                  ElasticSearch.builder()
                          .endpoints(server.getHttpHostAddress())
                          .build();
          client.connect();
  
          final String index = "test-index";
          assertTrue(client.index().create(index, null, null));
  
          final ImmutableMap<String, Object> doc = ImmutableMap.of("key", "val");
          final String idWithSpace = "an id"; // UI management templates' IDs contains spaces
          final String type = "type";
  
          client.documents().index(
              IndexRequest.builder()
                          .index(index)
                          .type(type)
                          .id(idWithSpace)
                          .doc(doc)
                          .build(), null);
  
          assertTrue(client.documents().get(index, type, idWithSpace).isPresent());
          assertEquals(client.documents().get(index, type, idWithSpace).get().getId(), idWithSpace);
          assertEquals(client.documents().get(index, type, idWithSpace).get().getSource(), doc);
  
          server.close();
      }
  
      @ParameterizedTest(name = "version: {0}")
      @MethodSource("es")
      public void testDocUpdate(final String ignored,
                                final ElasticsearchContainer server) {
          server.start();
  
          final ElasticSearch client =
                  ElasticSearch.builder()
                          .endpoints(server.getHttpHostAddress())
                          .build();
          client.connect();
  
          final String index = "test-index-update";
          assertTrue(client.index().create(index, null, null));
  
          final ImmutableMap<String, Object> doc = ImmutableMap.of("key", "val");
          final String idWithSpace = "an id"; // UI management templates' IDs contains spaces
          final String type = "type";
  
          client.documents().index(
              IndexRequest.builder()
                          .index(index)
                          .type(type)
                          .id(idWithSpace)
                          .doc(doc)
                          .build(), null);
  
          assertTrue(client.documents().get(index, type, idWithSpace).isPresent());
          assertEquals(client.documents().get(index, type, idWithSpace).get().getId(), idWithSpace);
          assertEquals(client.documents().get(index, type, idWithSpace).get().getSource(), doc);
  
          final Map<String, Object> updatedDoc = ImmutableMap.of("key", "new-val");
          client.documents().update(
              UpdateRequest
                  .builder()
                  .index(index)
                  .type(type)
                  .id(idWithSpace)
                  .doc(updatedDoc)
                  .build(),
              null);
          assertEquals(client.documents().get(index, type, idWithSpace).get().getSource(), updatedDoc);
  
          server.close();
      }
  
      @SuppressWarnings("unchecked")
      @ParameterizedTest(name = "version: {0}")
      @MethodSource("es")
      public void testSearch(final String ignored,
                             final ElasticsearchContainer server) {
          server.start();
  
          final ElasticSearch client =
                  ElasticSearch.builder()
                          .endpoints(server.getHttpHostAddress())
                          .build();
          client.connect();
  
          final String index = "test-index";
          final Mappings.Source sourceConf = new Mappings.Source();
          sourceConf.getExcludes().add("key3");
          assertTrue(
              client.index().create(
                  index,
                  Mappings.builder()
                          .type("type")
                          .properties(ImmutableMap.of("key1", ImmutableMap.of("type", "keyword")))
                          .properties(ImmutableMap.of("key2", ImmutableMap.of("type", "keyword"),
                                                      "key3", ImmutableMap.of("type", "keyword")
                          ))
                          .source(sourceConf)
                          .build(),
                  null
              )
          );
  
          final String type = "type";
  
          for (int i = 0; i < 10; i++) {
              client.documents().index(
                  IndexRequest.builder()
                              .index(index)
                              .type(type)
                              .id("id" + i)
                              .doc(ImmutableMap.of("key1", "val" + i, "key2", "val" + (i + 1), "key3", "val" + (i + 2)
                              ))
                              .build(), null);
          }
  
          await().atMost(Duration.ONE_MINUTE).untilAsserted(() -> {
              SearchResponse response = client.search(
                  Search.builder().query(Query.bool().must(Query.term("key1", "val1"))).build()
              );
              assertEquals(1, response.getHits().getTotal());
              assertEquals("val1", response.getHits().iterator().next().getSource().get("key1"));
          });
          //test indexOnly
          await().atMost(Duration.ONE_MINUTE).untilAsserted(() -> {
              SearchResponse response = client.search(
                  Search.builder().query(Query.bool().must(Query.term("key3", "val3"))).build()
              );
              assertEquals(1, response.getHits().getTotal());
              assertEquals("val1", response.getHits().iterator().next().getSource().get("key1"));
              assertNull(response.getHits().iterator().next().getSource().get("key3"), "indexOnly fields should not be stored");
          });
  
          await().atMost(Duration.ONE_MINUTE)
                 .pollInterval(Duration.FIVE_SECONDS)
                 .untilAsserted(() -> {
                     SearchResponse response = client.search(
                         Search.builder()
                               .query(
                                   Query.bool()
                                        .must(
                                            Query.bool()
                                                 .should(Query.bool()
                                                              .must(Query.term("key1", "val1"))
                                                              .must(Query.term("key2", "val2"))
                                                              .build())
                                                 .should(Query.bool()
                                                              .must(Query.term("key1", "val3"))
                                                              .must(Query.term("key2", "val4"))
                                                              .build()
                                                 )))
                               .aggregation(
                                   Aggregation
                                       .terms("key1").field("key1.keyword")
                                       .subAggregation(
                                           Aggregation.terms("key2").field("key2.keyword"))
                                       .size(1000)
                                       .build())
                               .build());
                     assertEquals(2, response.getHits().getTotal());
                     assertEquals(1, response.getAggregations().size());
                     assertEquals(
                         2,
                         (
                             (List<?>)
                                 ((Map<String, ?>) response.getAggregations().get("key1"))
                                     .get("buckets")
                         ).size()
                     );
  
                     //test mGet
                     Map<String, List<String>> indexIdsGroup = new HashMap<>();
                     indexIdsGroup.put("test-index", Arrays.asList("id1", "id2"));
                     Optional<Documents> documents = client.documents().mget(type, indexIdsGroup);
                     Map<String, Map<String, Object>> result = new HashMap<>();
                     assertThat(documents).isPresent();
                     for (final Document document : documents.get()) {
                         result.put(document.getId(), document.getSource());
                     }
                     assertEquals(2, result.get("id1").size());
                     assertEquals(2, result.get("id2").size());
                 });
  
          server.close();
      }
  }
  