Adds test

Signed-off-by: Gagan Juneja <[email protected]>
Gaganjuneja · May 16, 2024 · a32b4c9 · a32b4c9
1 parent 7a71dcd
commit a32b4c9
Showing 1 changed file with 66 additions and 0 deletions.
diff --git a/server/src/test/java/org/opensearch/node/NodeTests.java b/server/src/test/java/org/opensearch/node/NodeTests.java
@@ -34,17 +34,22 @@
 import org.apache.lucene.tests.util.LuceneTestCase;
 import org.opensearch.bootstrap.BootstrapCheck;
 import org.opensearch.bootstrap.BootstrapContext;
+import org.opensearch.client.Client;
 import org.opensearch.cluster.ClusterName;
+import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
 import org.opensearch.cluster.node.DiscoveryNodeRole;
+import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.SetOnce;
 import org.opensearch.common.network.NetworkModule;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.settings.SettingsException;
 import org.opensearch.core.common.breaker.CircuitBreaker;
+import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
 import org.opensearch.core.common.transport.BoundTransportAddress;
 import org.opensearch.core.common.unit.ByteSizeUnit;
 import org.opensearch.core.common.unit.ByteSizeValue;
 import org.opensearch.core.indices.breaker.CircuitBreakerService;
+import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.env.Environment;
 import org.opensearch.env.NodeEnvironment;
 import org.opensearch.index.IndexService;
@@ -56,22 +61,30 @@
 import org.opensearch.monitor.fs.FsProbe;
 import org.opensearch.plugins.CircuitBreakerPlugin;
 import org.opensearch.plugins.Plugin;
+import org.opensearch.plugins.TelemetryAwarePlugin;
+import org.opensearch.repositories.RepositoriesService;
+import org.opensearch.script.ScriptService;
+import org.opensearch.telemetry.metrics.MetricsRegistry;
+import org.opensearch.telemetry.tracing.Tracer;
 import org.opensearch.test.InternalTestCluster;
 import org.opensearch.test.MockHttpTransport;
 import org.opensearch.test.NodeRoles;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.watcher.ResourceWatcherService;
 
 import java.io.IOException;
 import java.nio.file.Path;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Supplier;
 
 import static org.opensearch.cluster.metadata.IndexMetadata.SETTING_NUMBER_OF_REPLICAS;
 import static org.opensearch.cluster.metadata.IndexMetadata.SETTING_NUMBER_OF_SHARDS;
@@ -404,6 +417,59 @@ public void testCreateWithFileCache() throws Exception {
         }
     }
 
+    public void testTelemetryAwarePlugins() throws IOException {
+        Settings.Builder settings = baseSettings();
+        List<Class<? extends Plugin>> plugins = basePlugins();
+        plugins.add(MockTelemetryAwarePlugin.class);
+        try (Node node = new MockNode(settings.build(), plugins)) {
+            MockTelemetryAwareComponent mockTelemetryAwareComponent = node.injector().getInstance(MockTelemetryAwareComponent.class);
+            assertNotNull(mockTelemetryAwareComponent.getTracer());
+            assertNotNull(mockTelemetryAwareComponent.getMetricsRegistry());
+            TelemetryAwarePlugin telemetryAwarePlugin = node.getPluginsService().filterPlugins(TelemetryAwarePlugin.class).get(0);
+            assertTrue(telemetryAwarePlugin instanceof MockTelemetryAwarePlugin);
+        }
+    }
+
+    private static class MockTelemetryAwareComponent {
+        private final Tracer tracer;
+        private final MetricsRegistry metricsRegistry;
+
+        public MockTelemetryAwareComponent(Tracer tracer, MetricsRegistry metricsRegistry) {
+            this.tracer = tracer;
+            this.metricsRegistry = metricsRegistry;
+        }
+
+        public Tracer getTracer() {
+            return tracer;
+        }
+
+        public MetricsRegistry getMetricsRegistry() {
+            return metricsRegistry;
+        }
+    }
+
+    public static class MockTelemetryAwarePlugin extends Plugin implements TelemetryAwarePlugin {
+        @Override
+        public Collection<Object> createComponents(
+            Client client,
+            ClusterService clusterService,
+            ThreadPool threadPool,
+            ResourceWatcherService resourceWatcherService,
+            ScriptService scriptService,
+            NamedXContentRegistry xContentRegistry,
+            Environment environment,
+            NodeEnvironment nodeEnvironment,
+            NamedWriteableRegistry namedWriteableRegistry,
+            IndexNameExpressionResolver indexNameExpressionResolver,
+            Supplier<RepositoriesService> repositoriesServiceSupplier,
+            Tracer tracer,
+            MetricsRegistry metricsRegistry
+        ) {
+            return List.of(new MockTelemetryAwareComponent(tracer, metricsRegistry));
+        }
+
+    }
+
     public static class MockCircuitBreakerPlugin extends Plugin implements CircuitBreakerPlugin {
 
         private SetOnce<CircuitBreaker> myCircuitBreaker = new SetOnce<>();