Added unit tests for diagnostics lib

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/listeners/RunnableEventListener.java

@@ -23,7 +23,7 @@ public interface RunnableEventListener {
      *              null span implies there is no active span
      * @param t    the thread executing the runnable
      */
-    void onRunnableStart(Span span, Thread t);
+    void onRunnableStart(final Span span, Thread t);
 
     /**
      * Called when a runnable completes execution on a given span
@@ -32,13 +32,13 @@ public interface RunnableEventListener {
      *             null span implies there is no active span
      * @param t    the thread executing the runnable
      */
-    void onRunnableComplete(Span span, Thread t);
+    void onRunnableComplete(final Span span, Thread t);
 
     /**
      * Checks whether the listener is enabled for the given span.
      * TODO - replace with operation based flag
      * @param span the span associated with the execution context
      * @return true if the listener is enabled for the span, false otherwise
      */
-    boolean isEnabled(Span span);
+    boolean isEnabled(final Span span);
 }

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/listeners/SpanEventListener.java

@@ -21,21 +21,21 @@ public interface SpanEventListener {
      * @param span the span that has started
      * @param t    the thread associated with the span
      */
-    void onSpanStart(Span span, Thread t);
+    void onSpanStart(final Span span, Thread t);
 
     /**
      * Called when a span completes.
      *
      * @param span the span that has completed
      * @param t    the thread associated with the span
      */
-    void onSpanComplete(Span span, Thread t);
+    void onSpanComplete(final Span span, Thread t);
 
     /**
      * Checks whether the listener is enabled for the given span.
      * TODO - replace with operation based flag
      * @param span the span to check
      * @return true if the listener is enabled for the span, false otherwise
      */
-    boolean isEnabled(Span span);
+    boolean isEnabled(final Span span);
 }

libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/listeners/TraceEventsService.java

@@ -35,7 +35,7 @@
  * The application must ensure that this service is updated with the latest set of TraceEventListener and latest value of {@link #tracingEnabled} and {@link #diagnosisEnabled}
  * are set, or it may produce undesirable results.
  */
-public final class TraceEventsService {
+public class TraceEventsService {
 
     private volatile Map<String, TraceEventListener> traceEventListeners;
     private volatile Tracer tracer;

libs/telemetry/src/test/java/org/opensearch/telemetry/diagnostics/DiagnosticsEventListenerTests.java

@@ -0,0 +1,99 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.telemetry.diagnostics;
+
+import org.mockito.ArgumentCaptor;
+import org.opensearch.telemetry.metrics.Measurement;
+import org.opensearch.telemetry.metrics.MetricEmitter;
+import org.opensearch.telemetry.metrics.MetricPoint;
+import org.opensearch.telemetry.tracing.Span;
+import org.opensearch.test.OpenSearchTestCase;
+
+import org.junit.Test;
+
+import java.util.Collections;
+
+import static org.mockito.Mockito.*;
+
+import org.mockito.Mockito;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.opensearch.telemetry.diagnostics.DiagnosticsEventListener.ELAPSED_TIME;
+
+public class DiagnosticsEventListenerTests extends OpenSearchTestCase {
+
+    private ThreadResourceRecorder<?> threadResourceRecorder;
+    private MetricEmitter metricEmitter;
+    private Span span;
+    private DiagnosticSpan diagnosticSpan;
+    private DiagnosticsEventListener diagnosticsEventListener;
+
+    @Override
+    public void setUp() throws Exception {
+        super.setUp();
+        threadResourceRecorder = Mockito.mock(ThreadResourceRecorder.class);
+        metricEmitter = Mockito.mock(MetricEmitter.class);
+        span = Mockito.mock(Span.class);
+        diagnosticSpan = Mockito.mock(DiagnosticSpan.class);
+        diagnosticsEventListener = new DiagnosticsEventListener(threadResourceRecorder, metricEmitter);
+    }
+
+    @Test
+    public void testOnSpanStart() {
+        Thread t = Thread.currentThread();
+        diagnosticsEventListener.onSpanStart(diagnosticSpan, t);
+        // Verify expected interactions
+        verify(threadResourceRecorder).startRecording(eq(diagnosticSpan), eq(t), eq(true));
+        verify(diagnosticSpan).putMetric(eq(DiagnosticsEventListener.START_SPAN_TIME), any(MetricPoint.class));
+    }
+
+    @Test
+    public void testOnSpanComplete() {
+        Thread t = Thread.currentThread();
+        MetricPoint diffMetric = new MetricPoint(Collections.emptyMap(), null, System.currentTimeMillis());
+        MetricPoint startMetric = new MetricPoint(Collections.emptyMap(), null, System.currentTimeMillis());
+        when(threadResourceRecorder.endRecording(any(DiagnosticSpan.class), eq(t), eq(true))).thenReturn(diffMetric);
+        when(diagnosticSpan.removeMetric(anyString())).thenReturn(startMetric);
+        ArgumentCaptor<MetricPoint> metricCaptor = ArgumentCaptor.forClass(MetricPoint.class);
+        diagnosticsEventListener.onSpanComplete(diagnosticSpan, t);
+        verify(metricEmitter).emitMetric(metricCaptor.capture());
+
+        // Check if diffMetric contains "elapsed_time" measurement
+        MetricPoint emittedMetric = metricCaptor.getValue();
+        Measurement<Number> elapsedTimeMeasurement = emittedMetric.getMeasurement(ELAPSED_TIME);
+        assertNotNull(elapsedTimeMeasurement);
+    }
+
+    @Test
+    public void testOnRunnableStart() {
+        Thread t = Thread.currentThread();
+        diagnosticsEventListener.onRunnableStart(diagnosticSpan, t);
+        verify(threadResourceRecorder).startRecording(eq(diagnosticSpan), eq(t), eq(false));
+    }
+
+    @Test
+    public void testOnRunnableComplete() {
+        Thread t = Thread.currentThread();
+        MetricPoint diffMetric = new MetricPoint(Collections.emptyMap(), null, System.currentTimeMillis());
+        when(threadResourceRecorder.endRecording(any(DiagnosticSpan.class), eq(t), eq(false))).thenReturn(diffMetric);
+
+        diagnosticsEventListener.onRunnableComplete(diagnosticSpan, t);
+
+        verify(metricEmitter).emitMetric(eq(diffMetric));
+    }
+
+    @Test
+    public void testIsEnabled() {
+        boolean isEnabled = diagnosticsEventListener.isEnabled(diagnosticSpan);
+        assertTrue(isEnabled);
+
+        isEnabled = diagnosticsEventListener.isEnabled(span);
+        assertFalse(isEnabled);
+    }
+}

libs/telemetry/src/test/java/org/opensearch/telemetry/diagnostics/ThreadResourceRecorderTests.java

@@ -0,0 +1,74 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.telemetry.diagnostics;
+
+import static org.mockito.Mockito.*;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.Collections;
+
+import org.opensearch.telemetry.metrics.MetricPoint;
+import org.opensearch.test.OpenSearchTestCase;
+
+public class ThreadResourceRecorderTests extends OpenSearchTestCase {
+
+    private ThreadResourceObserver observer;
+    private DiagnosticSpan span;
+    private Thread thread;
+
+    private static class TestThreadResourceRecorder extends ThreadResourceRecorder<ThreadResourceObserver> {
+        public TestThreadResourceRecorder(ThreadResourceObserver observer) {
+            super(observer);
+        }
+
+        @Override
+        protected MetricPoint computeDiff(MetricPoint startMetric, MetricPoint endMetric) {
+            // We simply return a new MetricPoint object with the same values as endMetric
+            return new MetricPoint(endMetric.getMeasurements(), null, endMetric.getObservationTime());
+        }
+    }
+
+    @Before
+    public void setUp() throws Exception {
+        super.setUp();
+
+        observer = mock(ThreadResourceObserver.class);
+        span = mock(DiagnosticSpan.class);
+        thread = mock(Thread.class);
+    }
+
+    @Test
+    public void testStartRecording() {
+        MetricPoint observation = new MetricPoint(Collections.emptyMap(), null, System.currentTimeMillis());
+        when(observer.observe(thread)).thenReturn(observation);
+
+        ThreadResourceRecorder<ThreadResourceObserver> recorder = new TestThreadResourceRecorder(observer);
+
+        recorder.startRecording(span, thread, true);
+
+        verify(span).putMetric(eq(String.valueOf(thread.getId())), eq(observation));
+    }
+
+    @Test
+    public void testEndRecording() {
+        MetricPoint startMetric = new MetricPoint(Collections.emptyMap(), null, System.currentTimeMillis());
+        MetricPoint endMetric = new MetricPoint(Collections.emptyMap(), null, System.currentTimeMillis() + 1000);
+        when(observer.observe(thread)).thenReturn(endMetric);
+
+        ThreadResourceRecorder<ThreadResourceObserver> recorder = new TestThreadResourceRecorder(observer);
+
+        when(span.removeMetric(String.valueOf(thread.getId()))).thenReturn(startMetric);
+
+        MetricPoint diffMetric = recorder.endRecording(span, thread, true);
+
+        assertEquals(endMetric.getObservationTime(), diffMetric.getObservationTime());
+    }
+}

libs/telemetry/src/test/java/org/opensearch/telemetry/tracing/listeners/TraceEventsRunnableTests.java

@@ -0,0 +1,117 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.telemetry.tracing.listeners;
+
+import org.junit.Test;
+import org.opensearch.telemetry.tracing.Span;
+import org.opensearch.telemetry.tracing.Tracer;
+import org.opensearch.test.OpenSearchTestCase;
+
+import static org.mockito.Mockito.*;
+
+public class TraceEventsRunnableTests extends OpenSearchTestCase {
+
+    private Tracer tracer;
+    private TraceEventsService traceEventsService;
+    private TraceEventListener traceEventListener1;
+    private TraceEventListener traceEventListener2;
+    private Span span;
+    private Thread currentThread;
+    private Runnable delegate;
+
+    @Override
+    public void setUp() throws Exception {
+        super.setUp();
+
+        tracer = mock(Tracer.class);
+        traceEventsService = spy(new TraceEventsService()); // Use spy here
+        traceEventListener1 = mock(TraceEventListener.class);
+        traceEventListener2 = mock(TraceEventListener.class);
+        span = mock(Span.class);
+        currentThread = mock(Thread.class);
+        delegate = mock(Runnable.class);
+
+        when(traceEventsService.getTracer()).thenReturn(tracer);
+        when(tracer.getCurrentSpan()).thenReturn(span);
+        when(span.getParentSpan()).thenReturn(null);
+
+        traceEventsService.registerTraceEventListener("listener1", traceEventListener1);
+        traceEventsService.registerTraceEventListener("listener2", traceEventListener2);
+        when(traceEventListener1.isEnabled(any(Span.class))).thenReturn(true);
+        when(traceEventListener2.isEnabled(any(Span.class))).thenReturn(true);
+
+        traceEventsService.setTracingEnabled(true);
+    }
+
+    @Test
+    public void testRun_InvokeOnRunnableStartAndOnRunnableComplete() {
+        Span span1 = mock(Span.class);
+        Span span2 = mock(Span.class);
+        when(traceEventsService.getTracer().getCurrentSpan()).thenReturn(span1, span1);
+        when(span1.hasEnded()).thenReturn(false);
+        when(span2.hasEnded()).thenReturn(false);
+        when(span1.getParentSpan()).thenReturn(span2);
+        when(span2.getParentSpan()).thenReturn(null);
+
+        TraceEventsRunnable traceEventsRunnable = new TraceEventsRunnable(delegate, traceEventsService);
+
+        traceEventsRunnable.run();
+
+        verify(traceEventListener1, times(2)).onRunnableStart(any(Span.class), any(Thread.class));
+        verify(traceEventListener2, times(2)).onRunnableStart(any(Span.class), any(Thread.class));
+        verify(traceEventListener1, times(2)).onRunnableComplete(any(Span.class), any(Thread.class));
+        verify(traceEventListener2, times(2)).onRunnableComplete(any(Span.class), any(Thread.class));
+
+        // Ensure that delegate.run() was invoked
+        verify(delegate).run();
+    }
+
+    @Test
+    public void testRun_TracingNotEnabled_NoInteractionsWithListeners() {
+        when(traceEventsService.isTracingEnabled()).thenReturn(false);
+
+        TraceEventsRunnable traceEventsRunnable = new TraceEventsRunnable(delegate, traceEventsService);
+
+        traceEventsRunnable.run();
+
+        // Verify that no interactions with listeners occurred
+        verifyNoInteractions(traceEventListener1);
+        verifyNoInteractions(traceEventListener2);
+    }
+
+    @Test
+    public void testRun_ExceptionInOnRunnableStart_NoImpactOnExecution() {
+        doThrow(new RuntimeException("Listener 1 exception")).when(traceEventListener1).onRunnableStart(eq(span), eq(currentThread));
+        TraceEventsRunnable traceEventsRunnable = new TraceEventsRunnable(delegate, traceEventsService);
+        traceEventsRunnable.run();
+
+        // Ensure that delegate.run() was invoked
+        verify(delegate).run();
+    }
+
+    @Test
+    public void testRun_ExceptionInOnRunnableComplete_NoImpactOnExecution() {
+        // trace event listener to throw an exception in onRunnableComplete
+        doThrow(new RuntimeException("Listener 1 exception")).when(traceEventListener1).onRunnableComplete(eq(span), eq(currentThread));
+        TraceEventsRunnable traceEventsRunnable = new TraceEventsRunnable(delegate, traceEventsService);
+        traceEventsRunnable.run();
+
+        // Verify that onRunnableStart was called for the listener despite the exception
+        verify(traceEventListener1).onRunnableStart(any(Span.class), any(Thread.class));
+        verify(delegate).run();
+    }
+
+    @Test
+    public void testUnwrap() {
+        TraceEventsRunnable traceEventsRunnable = new TraceEventsRunnable(delegate, traceEventsService);
+
+        Runnable unwrappedRunnable = traceEventsRunnable.unwrap();
+        assertSame(delegate, unwrappedRunnable);
+    }
+}