/*
 * Copyright 2021-2021 the original author or authors.
 *
 * Licensed 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
 *
 *      https://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.springframework.cloud.sleuth.instrument.messaging;

import java.util.function.Function;
import java.util.function.Supplier;

import org.junit.jupiter.api.Test;

import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.builder.SpringApplicationBuilder;
import org.springframework.cloud.function.context.FunctionCatalog;
import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry.FunctionInvocationWrapper;
import org.springframework.cloud.sleuth.test.TestSpanHandler;
import org.springframework.cloud.sleuth.test.TestTracer;
import org.springframework.cloud.sleuth.test.TestTracingBeanPostProcessor;
import org.springframework.context.ConfigurableApplicationContext;
import org.springframework.context.annotation.Bean;
import org.springframework.messaging.Message;
import org.springframework.messaging.support.MessageBuilder;

import static org.assertj.core.api.Assertions.assertThat;

/**
 * @author Oleg Zhurakousky
 *
 */
public abstract class TraceFunctionAroundWrapperTests {

	@Test
	public void test_tracing_with_supplier() {
		try (ConfigurableApplicationContext context = new SpringApplicationBuilder(configuration(),
				SampleConfiguration.class).run("--logging.level.org.springframework.cloud.function=DEBUG",
						"--spring.main.lazy-initialization=true");) {
			TestSpanHandler spanHandler = context.getBean(TestSpanHandler.class);
			assertThat(spanHandler.reportedSpans()).isEmpty();
			FunctionCatalog catalog = context.getBean(FunctionCatalog.class);
			FunctionInvocationWrapper function = catalog.lookup("greeter");

			Message<?> result = (Message<?>) function.get();

			assertThat(result.getPayload()).isEqualTo("hello".getBytes());
			assertThat(spanHandler.reportedSpans().size()).isEqualTo(2);
			assertThat(((String) result.getHeaders().get("b3"))).contains(spanHandler.get(0).getTraceId());
			spanHandler.assertAllSpansWereFinishedOrAbandoned(context.getBean(TestTracer.class).createdSpans());
		}
	}

	@Test
	public void test_tracing_with_function() {
		try (ConfigurableApplicationContext context = new SpringApplicationBuilder(configuration(),
				SampleConfiguration.class).run("--logging.level.org.springframework.cloud.function=DEBUG",
						"--spring.main.lazy-initialization=true");) {
			TestSpanHandler spanHandler = context.getBean(TestSpanHandler.class);
			assertThat(spanHandler.reportedSpans()).isEmpty();
			FunctionCatalog catalog = context.getBean(FunctionCatalog.class);
			FunctionInvocationWrapper function = catalog.lookup("uppercase");

			Message<?> result = (Message<?>) function.apply(MessageBuilder.withPayload("hello").build());

			assertThat(result.getPayload()).isEqualTo("HELLO");
			assertThat(spanHandler.reportedSpans().size()).isEqualTo(3);
			assertThat(((String) result.getHeaders().get("b3"))).contains(spanHandler.get(0).getTraceId());
			spanHandler.assertAllSpansWereFinishedOrAbandoned(context.getBean(TestTracer.class).createdSpans());
		}
	}

	protected abstract Class<?> configuration();

	@EnableAutoConfiguration
	public static class SampleConfiguration {

		@Bean
		public Supplier<String> greeter() {
			return () -> "hello";
		}

		@Bean
		public Function<String, String> uppercase() {
			return v -> v.toUpperCase();
		}

		@Bean
		static TestTracingBeanPostProcessor testTracerBeanPostProcessor() {
			return new TestTracingBeanPostProcessor();
		}

	}

};

@@ -787,16 +787,90 @@
 pi.Test;
+%0Aimport org.reactivestreams.Publisher;%0Aimport reactor.core.publisher.Flux;
 %0A%0Aimport
@@ -2701,16 +2701,989 @@
 %09%09%7D%0A%09%7D%0A%0A
+%09@Test%0A%09public void test_tracing_with_reactive_supplier() %7B%0A%09%09try (ConfigurableApplicationContext context = new SpringApplicationBuilder(configuration(),%0A%09%09%09%09SampleConfiguration.class).run(%22--logging.level.org.springframework.cloud.function=DEBUG%22,%0A%09%09%09%09%09%09%22--spring.main.lazy-initialization=true%22);) %7B%0A%09%09%09TestSpanHandler spanHandler = context.getBean(TestSpanHandler.class);%0A%09%09%09assertThat(spanHandler.reportedSpans()).isEmpty();%0A%09%09%09FunctionCatalog catalog = context.getBean(FunctionCatalog.class);%0A%09%09%09FunctionInvocationWrapper function = catalog.lookup(%22reactiveGreeter%22);%0A%09%09%09function.setSkipOutputConversion(true);%0A%09%09%09Object result = function.get();%0A%09%09%09assertThat(result).isInstanceOf(Publisher.class);%0A%09%09%09/* TODO%0A%09%09%09 * We'll need more assertions but for now this one will ensure that wrapper does not change the type of return value%0A%09%09%09 * specifically for reactive cases where Flux became Message%3CFlux%3E due to the current code in TraceFunctionAroundWrapper%0A%09%09%09 */%0A%09%09%7D%0A%09%7D%0A%0A
 %09@Test%0A%09
@@ -4819,24 +4819,24 @@
 greeter() %7B%0A
-
 %09%09%09return ()
@@ -4849,24 +4849,125 @@
 ello%22;%0A%09%09%7D%0A%0A
+%09%09@Bean%0A%09%09public Supplier%3CFlux%3CString%3E%3E reactiveGreeter() %7B%0A%09%09%09return () -%3E Flux.just(%22hello%22);%0A%09%09%7D%0A%0A
 %09%09@Bean%0A%09%09pu

Add reactive Supplier test