From 25606b289622a246cfa4569e6d4e5b07bc782f96 Mon Sep 17 00:00:00 2001
From: maciejtoporowicz <maciej.toporowicz@gmail.com>
Date: Tue, 14 Mar 2023 17:16:29 +0100
Subject: [PATCH] Add getters for delegates (#2273)

Fixes gh-2261
---
 .../sleuth/instrument/async/TraceCallable.java |  7 +++++++
 .../sleuth/instrument/async/TraceRunnable.java |  7 +++++++
 .../instrument/async/TraceCallableTests.java   | 18 ++++++++++++++++++
 .../instrument/async/TraceRunnableTests.java   | 16 ++++++++++++++++
 4 files changed, 48 insertions(+)

diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceCallable.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceCallable.java
index 2aeb420f3..b2d08d24a 100644
--- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceCallable.java
+++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceCallable.java
@@ -75,4 +75,11 @@ public class TraceCallable<V> implements Callable<V> {
 		}
 	}
 
+	/**
+	 * @return delegate {@link Callable}
+	 */
+	public Callable<V> getDelegate() {
+		return delegate;
+	}
+
 }
diff --git a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceRunnable.java b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceRunnable.java
index de95a19e2..bb886e86b 100644
--- a/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceRunnable.java
+++ b/spring-cloud-sleuth-instrumentation/src/main/java/org/springframework/cloud/sleuth/instrument/async/TraceRunnable.java
@@ -72,4 +72,11 @@ public class TraceRunnable implements Runnable {
 		}
 	}
 
+	/**
+	 * @return delegate {@link Runnable}
+	 */
+	public Runnable getDelegate() {
+		return delegate;
+	}
+
 }
diff --git a/tests/brave/spring-cloud-sleuth-instrumentation-async-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/async/TraceCallableTests.java b/tests/brave/spring-cloud-sleuth-instrumentation-async-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/async/TraceCallableTests.java
index 1dd6a43f0..8cf2db7a1 100644
--- a/tests/brave/spring-cloud-sleuth-instrumentation-async-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/async/TraceCallableTests.java
+++ b/tests/brave/spring-cloud-sleuth-instrumentation-async-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/async/TraceCallableTests.java
@@ -16,9 +16,17 @@
 
 package org.springframework.cloud.sleuth.brave.instrument.async;
 
+import java.util.concurrent.Callable;
+
+import org.junit.jupiter.api.Test;
+
 import org.springframework.cloud.sleuth.brave.BraveTestTracing;
+import org.springframework.cloud.sleuth.instrument.async.TraceCallable;
+import org.springframework.cloud.sleuth.internal.DefaultSpanNamer;
 import org.springframework.cloud.sleuth.test.TestTracingAware;
 
+import static org.assertj.core.api.Assertions.assertThat;
+
 public class TraceCallableTests extends org.springframework.cloud.sleuth.instrument.async.TraceCallableTests {
 
 	BraveTestTracing testTracing;
@@ -31,4 +39,14 @@ public class TraceCallableTests extends org.springframework.cloud.sleuth.instrum
 		return this.testTracing;
 	}
 
+	@Test
+	public void should_provide_delegate() {
+		Callable<String> delegate = () -> "test";
+
+		TraceCallable<String> traceCallable = new TraceCallable<>(tracerTest().tracing().tracer(),
+				new DefaultSpanNamer(), delegate);
+
+		assertThat(traceCallable.getDelegate()).isEqualTo(delegate);
+	}
+
 }
diff --git a/tests/brave/spring-cloud-sleuth-instrumentation-async-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/async/TraceRunnableTests.java b/tests/brave/spring-cloud-sleuth-instrumentation-async-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/async/TraceRunnableTests.java
index 4b32d672e..c9373c3d1 100644
--- a/tests/brave/spring-cloud-sleuth-instrumentation-async-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/async/TraceRunnableTests.java
+++ b/tests/brave/spring-cloud-sleuth-instrumentation-async-tests/src/test/java/org/springframework/cloud/sleuth/brave/instrument/async/TraceRunnableTests.java
@@ -17,11 +17,16 @@
 package org.springframework.cloud.sleuth.brave.instrument.async;
 
 import org.assertj.core.api.BDDAssertions;
+import org.junit.jupiter.api.Test;
 
 import org.springframework.cloud.sleuth.Span;
 import org.springframework.cloud.sleuth.brave.BraveTestTracing;
+import org.springframework.cloud.sleuth.instrument.async.TraceRunnable;
+import org.springframework.cloud.sleuth.internal.DefaultSpanNamer;
 import org.springframework.cloud.sleuth.test.TestTracingAware;
 
+import static org.assertj.core.api.Assertions.assertThat;
+
 public class TraceRunnableTests extends org.springframework.cloud.sleuth.instrument.async.TraceRunnableTests {
 
 	BraveTestTracing testTracing;
@@ -39,4 +44,15 @@ public class TraceRunnableTests extends org.springframework.cloud.sleuth.instrum
 		BDDAssertions.then(secondSpan.context().parentId()).as("saved span as remnant of first span").isNull();
 	}
 
+	@Test
+	public void should_provide_delegate() {
+		Runnable delegate = () -> {
+		};
+
+		TraceRunnable traceRunnable = new TraceRunnable(tracerTest().tracing().tracer(), new DefaultSpanNamer(),
+				delegate);
+
+		assertThat(traceRunnable.getDelegate()).isEqualTo(delegate);
+	}
+
 }