Ability to extend default function result handling

Fixes #757

spring-cloud-function-adapters/spring-cloud-function-adapter-azure/src/test/java/org/springframework/cloud/function/adapter/azure/CustomFunctionInvokerTests.java

@@ -0,0 +1,200 @@
+/*
+ * Copyright 2022-2022 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.function.adapter.azure;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+import com.microsoft.azure.functions.ExecutionContext;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import org.reactivestreams.Publisher;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry.FunctionInvocationWrapper;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.messaging.support.GenericMessage;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.util.Lists.list;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyList;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.same;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+
+/**
+ * Unit tests for {@link FunctionInvoker} custom result handling.
+ *
+ * @author Chris Bono
+ */
+class CustomFunctionInvokerTests {
+
+	private FunctionInvoker<?, ?>  currentInvoker;
+
+	@AfterEach
+	void closeCurrentInvoker() {
+		if (this.currentInvoker != null) {
+			this.currentInvoker.close();
+		}
+	}
+
+	/**
+	 * Verifies custom result handling and proper post-process callback invocation for an imperative function.
+	 */
+	@Test
+	void customImperativeResultHandling() {
+		FunctionInvoker<String, String> invoker = new FunctionInvoker<String, String>(TestFunctionsConfig.class) {
+			@Override
+			protected String postProcessImperativeFunctionResult(String rawInputs, Object functionInputs,
+				Object functionResult, FunctionInvocationWrapper function, ExecutionContext executionContext
+			) {
+				return functionResult + "+imperative";
+			}
+		};
+		invoker = spyOnAndCloseAfterTest(invoker);
+		ExecutionContext executionContext = new TestExecutionContext("imperativeUppercase");
+		String result = invoker.handleRequest("foo", executionContext);
+		assertThat(result).isEqualTo("FOO+imperative");
+
+		// Below here verifies that the expected callback(s) were invoked w/ the expected arguments
+
+		// Only imperative post-process callback should be called
+		verify(invoker, never()).postProcessReactiveFunctionResult(anyString(), any(), any(Publisher.class), any(), same(executionContext));
+		verify(invoker, never()).postProcessMonoFunctionResult(anyString(), any(), any(Mono.class), any(), same(executionContext));
+		verify(invoker, never()).postProcessFluxFunctionResult(anyString(), any(), any(Flux.class), any(), same(executionContext));
+
+		// Only sniff-test the payload of the input message (the other fields are problematic to verify and no value doing that here)
+		ArgumentCaptor<GenericMessage> functionInputsCaptor = ArgumentCaptor.forClass(GenericMessage.class);
+		verify(invoker).postProcessImperativeFunctionResult(eq("foo"), functionInputsCaptor.capture(), eq("FOO"), any(), same(executionContext));
+		assertThat(functionInputsCaptor.getValue()).extracting(GenericMessage::getPayload).isEqualTo("foo");
+	}
+
+	/**
+	 * Verifies custom result handling and proper post-process callback invocation for a reactive Mono function.
+	 */
+	@Test
+	void customReactiveMonoResultHandling() {
+		FunctionInvoker<String, String> invoker = new FunctionInvoker<String, String>(TestFunctionsConfig.class) {
+			@Override
+			protected String postProcessMonoFunctionResult(String rawInputs, Object functionInputs, Mono<?> functionResult,
+				FunctionInvocationWrapper function, ExecutionContext executionContext
+			) {
+				return functionResult.block().toString() + "+mono";
+			}
+		};
+		invoker = spyOnAndCloseAfterTest(invoker);
+		ExecutionContext executionContext = new TestExecutionContext("reactiveMonoUppercase");
+		String result = invoker.handleRequest("foo", executionContext);
+		assertThat(result).isEqualTo("FOO+mono");
+
+		// Below here verifies that the expected callback(s) were invoked w/ the expected arguments
+
+		// Only publisher->mono post-process callbacks should be called
+		verify(invoker, never()).postProcessImperativeFunctionResult(anyString(), any(), any(), any(), same(executionContext));
+		verify(invoker, never()).postProcessFluxFunctionResult(anyString(), any(), any(Flux.class), any(), same(executionContext));
+
+		// Only sniff-test the payload of the input message and the mono (the other fields are problematic to verify and no value doing that here)
+		ArgumentCaptor<GenericMessage> functionInputsCaptor = ArgumentCaptor.forClass(GenericMessage.class);
+		ArgumentCaptor<Mono> functionResultCaptor = ArgumentCaptor.forClass(Mono.class);
+		verify(invoker).postProcessReactiveFunctionResult(eq("foo"), functionInputsCaptor.capture(), functionResultCaptor.capture(), any(), same(executionContext));
+		verify(invoker).postProcessMonoFunctionResult(eq("foo"), functionInputsCaptor.capture(), functionResultCaptor.capture(), any(), same(executionContext));
+		// NOTE: The captors get called twice as the args are just delegated from publisher->mono callback
+		assertThat(functionInputsCaptor.getAllValues()).extracting(GenericMessage::getPayload).containsExactly("foo", "foo");
+		assertThat(functionResultCaptor.getAllValues()).extracting(Mono::block).containsExactly("FOO", "FOO");
+	}
+
+	/**
+	 * Verifies custom result handling and proper post-process callback invocation for a reactive Flux function.
+	 */
+	@Test
+	void customReactiveFluxResultHandling() {
+		FunctionInvoker<List<String>, String> invoker = new FunctionInvoker<List<String>, String>(TestFunctionsConfig.class) {
+			@Override
+			protected String postProcessFluxFunctionResult(List<String> rawInputs, Object functionInputs,
+				Flux<?> functionResult, FunctionInvocationWrapper function, ExecutionContext executionContext
+			) {
+				return functionResult.map(o -> o.toString() + "+flux").collectList().block().stream().collect(Collectors.joining("/"));
+			}
+		};
+		invoker = spyOnAndCloseAfterTest(invoker);
+		ExecutionContext executionContext = new TestExecutionContext("reactiveFluxUppercase");
+		List<String> rawInputs = Arrays.asList("foo", "bar");
+		String result = invoker.handleRequest(rawInputs, executionContext);
+		assertThat(result).isEqualTo("FOO+flux/BAR+flux");
+
+		// Below here verifies that the expected callback(s) were invoked w/ the expected arguments
+
+		// Only publisher->flux post-process callbacks should be called
+		verify(invoker, never()).postProcessImperativeFunctionResult(anyList(), any(), any(), any(), same(executionContext));
+		verify(invoker, never()).postProcessMonoFunctionResult(anyList(), any(), any(Mono.class), any(), same(executionContext));
+
+		// Only sniff-test the payload of the input message and the mono (the other fields are problematic to verify and no value doing that here)
+		ArgumentCaptor<Flux<GenericMessage>> functionInputsCaptor = ArgumentCaptor.forClass(Flux.class);
+		ArgumentCaptor<Flux> functionResultCaptor = ArgumentCaptor.forClass(Flux.class);
+		verify(invoker).postProcessReactiveFunctionResult(same(rawInputs), functionInputsCaptor.capture(), functionResultCaptor.capture(), any(), same(executionContext));
+		verify(invoker).postProcessFluxFunctionResult(same(rawInputs), functionInputsCaptor.capture(), functionResultCaptor.capture(), any(), same(executionContext));
+
+		// NOTE: The captors get called twice as the args are just delegated from publisher->flux callback
+
+		// The functionInputs for each call is Flux<GreetingMessage> with 2 items - one for 'foo' and one for 'bar'
+		assertThat(functionInputsCaptor.getAllValues())
+			.extracting(Flux::collectList).extracting(Mono::block)
+			.flatExtracting(fluxAsList -> fluxAsList.stream().collect(Collectors.toList()))
+			.extracting(GenericMessage::getPayload).containsExactlyInAnyOrder("foo", "bar", "foo", "bar");
+
+		// The functionResult for each call is a Flux<String> w/ 2 items { "FOO", "BAR" }
+		assertThat(functionResultCaptor.getAllValues())
+			.extracting(Flux::collectList).extracting(Mono::block)
+			.containsExactlyInAnyOrder(list("FOO", "BAR"), list("FOO", "BAR"));
+	}
+
+	private <I, O> FunctionInvoker<I, O> spyOnAndCloseAfterTest(FunctionInvoker<I, O> invoker) {
+		this.currentInvoker = invoker;
+		return spy(invoker);
+	}
+
+	@Configuration
+	@EnableAutoConfiguration
+	static class TestFunctionsConfig {
+
+		@Bean
+		public Function<String, String> imperativeUppercase() {
+			return (s) -> s.toUpperCase();
+		}
+
+		@Bean
+		public Function<Mono<String>, Mono<String>> reactiveMonoUppercase() {
+			return (m) -> m.map(String::toUpperCase);
+		}
+
+		@Bean
+		public Function<Flux<String>, Flux<String>> reactiveFluxUppercase() {
+			return (f) -> f.map(String::toUpperCase);
+		}
+
+	}
+}