More tests

Author: smalyshev

server/src/internalClusterTest/java/org/elasticsearch/search/ccs/CCSUsageTelemetryIT.java

@@ -40,27 +40,19 @@
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.test.AbstractMultiClustersTestCase;
 import org.elasticsearch.test.InternalTestCluster;
+import org.elasticsearch.test.SkipUnavailableRule;
+import org.elasticsearch.test.SkipUnavailableRule.NotSkipped;
 import org.elasticsearch.usage.UsageService;
 import org.junit.Assert;
 import org.junit.Rule;
-import org.junit.rules.TestRule;
-import org.junit.runner.Description;
-import org.junit.runners.model.Statement;
-
-import java.lang.annotation.ElementType;
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
-import java.lang.annotation.Target;
-import java.util.Arrays;
+
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
-import java.util.function.Function;
-import java.util.stream.Collectors;
 
 import static org.elasticsearch.action.admin.cluster.stats.CCSUsageTelemetry.ASYNC_FEATURE;
 import static org.elasticsearch.action.admin.cluster.stats.CCSUsageTelemetry.MRT_FEATURE;
@@ -498,7 +490,7 @@ public void testRemoteOnlyTimesOut() throws Exception {
         assertThat(perCluster.get(REMOTE2), equalTo(null));
     }
 
-    @SkipOverride(aliases = { REMOTE1 })
+    @NotSkipped(aliases = { REMOTE1 })
     public void testRemoteTimesOutFailure() throws Exception {
         Map<String, Object> testClusterInfo = setupClusters();
         String remoteIndex = (String) testClusterInfo.get("remote.index");
@@ -528,7 +520,7 @@ public void testRemoteTimesOutFailure() throws Exception {
     /**
     * Search when all the remotes failed and not skipped
     */
-    @SkipOverride(aliases = { REMOTE1, REMOTE2 })
+    @NotSkipped(aliases = { REMOTE1, REMOTE2 })
     public void testFailedAllRemotesSearch() throws Exception {
         Map<String, Object> testClusterInfo = setupClusters();
         String localIndex = (String) testClusterInfo.get("local.index");
@@ -577,7 +569,7 @@ public void testRemoteHasNoIndex() throws Exception {
     /**
      * Test that we're still counting remote search even if remote cluster has no such index
      */
-    @SkipOverride(aliases = { REMOTE1 })
+    @NotSkipped(aliases = { REMOTE1 })
     public void testRemoteHasNoIndexFailure() throws Exception {
         SearchRequest searchRequest = makeSearchRequest(REMOTE1 + ":no_such_index");
         CCSTelemetrySnapshot telemetry = getTelemetryFromFailedSearch(searchRequest);
@@ -695,40 +687,4 @@ private void indexDocs(Client client, String index, ActionListener<Void> listene
         bulkRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE).execute(listener.safeMap(r -> null));
     }
 
-    /**
-     * Annotation to mark specific cluster in a test as not to be skipped when unavailable
-     */
-    @Retention(RetentionPolicy.RUNTIME)
-    @Target(ElementType.METHOD)
-    @interface SkipOverride {
-        String[] aliases();
-    }
-
-    /**
-     * Test rule to process skip annotations
-     */
-    static class SkipUnavailableRule implements TestRule {
-        private final Map<String, Boolean> skipMap;
-
-        SkipUnavailableRule(String... clusterAliases) {
-            this.skipMap = Arrays.stream(clusterAliases).collect(Collectors.toMap(Function.identity(), alias -> true));
-        }
-
-        public Map<String, Boolean> getMap() {
-            return skipMap;
-        }
-
-        @Override
-        public Statement apply(Statement base, Description description) {
-            // Check for annotation named "SkipOverride" and set the overrides accordingly
-            var aliases = description.getAnnotation(SkipOverride.class);
-            if (aliases != null) {
-                for (String alias : aliases.aliases()) {
-                    skipMap.put(alias, false);
-                }
-            }
-            return base;
-        }
-
-    }
 }

server/src/main/java/org/elasticsearch/action/admin/cluster/stats/CCSUsage.java

@@ -10,6 +10,7 @@
 package org.elasticsearch.action.admin.cluster.stats;
 
 import org.elasticsearch.ElasticsearchSecurityException;
+import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.ResourceNotFoundException;
 import org.elasticsearch.action.ShardOperationFailedException;
@@ -143,6 +144,9 @@ public static Result getFailureType(Exception e) {
             if (ExceptionsHelper.unwrapCorruption(e) != null) {
                 return Result.CORRUPTION;
             }
+            if (ExceptionsHelper.unwrap(e, ElasticsearchStatusException.class) != null) {
+                return Result.LICENSE;
+            }
             // This is kind of last resort check - if we still don't know the reason but all shard failures are remote,
             // we assume it's remote's fault somehow.
             if (e instanceof SearchPhaseExecutionException spe) {

server/src/main/java/org/elasticsearch/action/admin/cluster/stats/CCSUsageTelemetry.java

@@ -47,6 +47,7 @@ public enum Result {
         TIMEOUT("timeout"),
         CORRUPTION("corruption"),
         SECURITY("security"),
+        LICENSE("license"),
         // May be helpful if there's a lot of other reasons, and it may be hard to calculate the unknowns for some clients.
         UNKNOWN("other");
 

test/framework/src/main/java/org/elasticsearch/test/SkipUnavailableRule.java

@@ -0,0 +1,60 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.test;
+
+import org.junit.rules.TestRule;
+import org.junit.runner.Description;
+import org.junit.runners.model.Statement;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+import java.util.Arrays;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+/**
+ * Test rule to process skip_unavailable override annotations
+ */
+public class SkipUnavailableRule implements TestRule {
+    private final Map<String, Boolean> skipMap;
+
+    public SkipUnavailableRule(String... clusterAliases) {
+        this.skipMap = Arrays.stream(clusterAliases).collect(Collectors.toMap(Function.identity(), alias -> true));
+    }
+
+    public Map<String, Boolean> getMap() {
+        return skipMap;
+    }
+
+    @Override
+    public Statement apply(Statement base, Description description) {
+        // Check for annotation named "SkipOverride" and set the overrides accordingly
+        var aliases = description.getAnnotation(NotSkipped.class);
+        if (aliases != null) {
+            for (String alias : aliases.aliases()) {
+                skipMap.put(alias, false);
+            }
+        }
+        return base;
+    }
+
+    /**
+     * Annotation to mark specific cluster in a test as not to be skipped when unavailable
+     */
+    @Retention(RetentionPolicy.RUNTIME)
+    @Target(ElementType.METHOD)
+    public @interface NotSkipped {
+        String[] aliases();
+    }
+
+}

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/CrossClustersUsageTelemetryIT.java

@@ -11,34 +11,21 @@
 import org.apache.logging.log4j.Logger;
 import org.elasticsearch.action.admin.cluster.stats.CCSTelemetrySnapshot;
 import org.elasticsearch.client.internal.Client;
-import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Settings;
-import org.elasticsearch.compute.operator.exchange.ExchangeService;
-import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.test.AbstractMultiClustersTestCase;
+import org.elasticsearch.test.SkipUnavailableRule;
 import org.elasticsearch.usage.UsageService;
-import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
 import org.junit.Assert;
 import org.junit.Rule;
-import org.junit.rules.TestRule;
-import org.junit.runner.Description;
-import org.junit.runners.model.Statement;
-
-import java.lang.annotation.ElementType;
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
-import java.lang.annotation.Target;
+
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ExecutionException;
-import java.util.function.Function;
-import java.util.stream.Collectors;
 
 import static org.elasticsearch.action.admin.cluster.stats.CCSUsageTelemetry.ASYNC_FEATURE;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
@@ -52,19 +39,6 @@ public class CrossClustersUsageTelemetryIT extends AbstractMultiClustersTestCase
     private static final String LOCAL_INDEX = "logs-1";
     private static final String REMOTE_INDEX = "logs-2";
 
-    @Override
-    protected boolean reuseClusters() {
-        return false;
-    }
-
-    @Override
-    protected List<String> remoteClusterAlias() {
-        return List.of(REMOTE1, REMOTE2);
-    }
-
-    @Rule
-    public SkipUnavailableRule skipOverride = new SkipUnavailableRule(REMOTE1, REMOTE2);
-
     public void testLocalRemote() throws Exception {
         setupClusters();
         var telemetry = getTelemetryFromQuery("from logs-*,c*:logs-* | stats sum (v)", "kibana");
@@ -93,11 +67,7 @@ public void testLocalRemote() throws Exception {
 
     }
 
-    private CCSTelemetrySnapshot getTelemetryFromQuery(String query) throws ExecutionException, InterruptedException {
-        return getTelemetryFromQuery(query, null);
-    }
-
-    private CCSTelemetrySnapshot getTelemetryFromQuery(String query, String client) throws ExecutionException, InterruptedException {
+    protected CCSTelemetrySnapshot getTelemetryFromQuery(String query, String client) throws ExecutionException, InterruptedException {
         EsqlQueryRequest request = EsqlQueryRequest.syncEsqlQueryRequest();
         request.query(query);
         request.pragmas(AbstractEsqlIntegTestCase.randomPragmas());
@@ -106,7 +76,7 @@ private CCSTelemetrySnapshot getTelemetryFromQuery(String query, String client)
         return getTelemetryFromQuery(request, client);
     }
 
-    private CCSTelemetrySnapshot getTelemetryFromQuery(EsqlQueryRequest request, String client) throws ExecutionException,
+    protected CCSTelemetrySnapshot getTelemetryFromQuery(EsqlQueryRequest request, String client) throws ExecutionException,
         InterruptedException {
         // We want to send search to a specific node (we don't care which one) so that we could
         // collect the CCS telemetry from it later
@@ -116,7 +86,7 @@ private CCSTelemetrySnapshot getTelemetryFromQuery(EsqlQueryRequest request, Str
         if (client != null) {
             assertResponse(
                 cluster(LOCAL_CLUSTER).client(nodeName)
-                    .filterWithHeader(Map.of(Task.X_ELASTIC_PRODUCT_ORIGIN_HTTP_HEADER, "kibana"))
+                    .filterWithHeader(Map.of(Task.X_ELASTIC_PRODUCT_ORIGIN_HTTP_HEADER, client))
                     .execute(EsqlQueryAction.INSTANCE, request),
                 Assert::assertNotNull
             );
@@ -127,12 +97,44 @@ private CCSTelemetrySnapshot getTelemetryFromQuery(EsqlQueryRequest request, Str
         return getTelemetrySnapshot(nodeName);
     }
 
+    protected CCSTelemetrySnapshot getTelemetryFromFailedQuery(String query) throws Exception {
+        // We want to send search to a specific node (we don't care which one) so that we could
+        // collect the CCS telemetry from it later
+        String nodeName = cluster(LOCAL_CLUSTER).getRandomNodeName();
+        EsqlQueryRequest request = EsqlQueryRequest.syncEsqlQueryRequest();
+        request.query(query);
+        request.pragmas(AbstractEsqlIntegTestCase.randomPragmas());
+        request.columnar(randomBoolean());
+        request.includeCCSMetadata(randomBoolean());
+
+        ExecutionException ee = expectThrows(
+            ExecutionException.class,
+            cluster(LOCAL_CLUSTER).client(nodeName).execute(EsqlQueryAction.INSTANCE, request)::get
+        );
+        assertNotNull(ee.getCause());
+
+        return getTelemetrySnapshot(nodeName);
+    }
+
     private CCSTelemetrySnapshot getTelemetrySnapshot(String nodeName) {
         var usage = cluster(LOCAL_CLUSTER).getInstance(UsageService.class, nodeName);
         return usage.getEsqlUsageHolder().getCCSTelemetrySnapshot();
     }
 
-    Map<String, Object> setupClusters() {
+    @Override
+    protected boolean reuseClusters() {
+        return false;
+    }
+
+    @Override
+    protected List<String> remoteClusterAlias() {
+        return List.of(REMOTE1, REMOTE2);
+    }
+
+    @Rule
+    public SkipUnavailableRule skipOverride = new SkipUnavailableRule(REMOTE1, REMOTE2);
+
+    protected Map<String, Object> setupClusters() {
         int numShardsLocal = randomIntBetween(1, 5);
         populateLocalIndices(LOCAL_INDEX, numShardsLocal);
 
@@ -186,65 +188,15 @@ void populateRemoteIndices(String clusterAlias, String indexName, int numShards)
     @Override
     protected Collection<Class<? extends Plugin>> nodePlugins(String clusterAlias) {
         List<Class<? extends Plugin>> plugins = new ArrayList<>(super.nodePlugins(clusterAlias));
-        plugins.add(EsqlPlugin.class);
+        plugins.add(EsqlPluginWithEnterpriseOrTrialLicense.class);
         plugins.add(CrossClustersQueryIT.InternalExchangePlugin.class);
         return plugins;
     }
 
-    public static class InternalExchangePlugin extends Plugin {
-        @Override
-        public List<Setting<?>> getSettings() {
-            return List.of(
-                Setting.timeSetting(
-                    ExchangeService.INACTIVE_SINKS_INTERVAL_SETTING,
-                    TimeValue.timeValueSeconds(30),
-                    Setting.Property.NodeScope
-                )
-            );
-        }
-    }
-
     @Override
     protected Map<String, Boolean> skipUnavailableForRemoteClusters() {
         var map = skipOverride.getMap();
         LOGGER.info("Using skip_unavailable map: [{}]", map);
         return map;
     }
-
-    /**
-     * Annotation to mark specific cluster in a test as not to be skipped when unavailable
-     */
-    @Retention(RetentionPolicy.RUNTIME)
-    @Target(ElementType.METHOD)
-    @interface SkipOverride {
-        String[] aliases();
-    }
-
-    /**
-     * Test rule to process skip annotations
-     */
-    static class SkipUnavailableRule implements TestRule {
-        private final Map<String, Boolean> skipMap;
-
-        SkipUnavailableRule(String... clusterAliases) {
-            this.skipMap = Arrays.stream(clusterAliases).collect(Collectors.toMap(Function.identity(), alias -> true));
-        }
-
-        public Map<String, Boolean> getMap() {
-            return skipMap;
-        }
-
-        @Override
-        public Statement apply(Statement base, Description description) {
-            // Check for annotation named "SkipOverride" and set the overrides accordingly
-            var aliases = description.getAnnotation(SkipOverride.class);
-            if (aliases != null) {
-                for (String alias : aliases.aliases()) {
-                    skipMap.put(alias, false);
-                }
-            }
-            return base;
-        }
-
-    }
 }