test: Add unit test for MessageAggregator tool call aggregation

Author: WOONBE

spring-ai-model/src/test/java/org/springframework/ai/chat/model/ChatResponseTests.java

@@ -19,14 +19,18 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicReference;
 
 import org.junit.jupiter.api.Test;
 
 import org.springframework.ai.chat.messages.AssistantMessage;
 import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
+import org.springframework.ai.chat.metadata.ChatResponseMetadata;
+import reactor.core.publisher.Flux;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.springframework.ai.chat.messages.AssistantMessage.*;
 
 /**
  * Unit tests for {@link ChatResponse}.
@@ -38,8 +42,8 @@ class ChatResponseTests {
 	@Test
 	void whenToolCallsArePresentThenReturnTrue() {
 		ChatResponse chatResponse = ChatResponse.builder()
-			.generations(List.of(new Generation(new AssistantMessage("", Map.of(),
-					List.of(new AssistantMessage.ToolCall("toolA", "function", "toolA", "{}"))))))
+			.generations(List.of(new Generation(
+					new AssistantMessage("", Map.of(), List.of(new ToolCall("toolA", "function", "toolA", "{}"))))))
 			.build();
 		assertThat(chatResponse.hasToolCalls()).isTrue();
 	}
@@ -80,4 +84,45 @@ void whenFinishReasonIsNotPresent() {
 		assertThat(chatResponse.hasFinishReasons(Set.of("completed"))).isFalse();
 	}
 
+	@Test
+	void messageAggregatorShouldCorrectlyAggregateToolCallsFromStream() {
+
+		MessageAggregator aggregator = new MessageAggregator();
+
+		ChatResponse chunk1 = new ChatResponse(
+				List.of(new Generation(new AssistantMessage("Thinking about the weather... "))));
+
+		ToolCall weatherToolCall = new ToolCall("tool-id-123", "function", "getCurrentWeather",
+				"{\"location\": \"Seoul\"}");
+
+		Map<String, Object> metadataWithToolCall = Map.of("toolCalls", List.of(weatherToolCall));
+		ChatResponseMetadata responseMetadataForChunk2 = ChatResponseMetadata.builder()
+			.metadata(metadataWithToolCall)
+			.build();
+
+		ChatResponse chunk2 = new ChatResponse(List.of(new Generation(new AssistantMessage(""))),
+				responseMetadataForChunk2);
+
+		Flux<ChatResponse> streamingResponse = Flux.just(chunk1, chunk2);
+
+		AtomicReference<ChatResponse> aggregatedResponseRef = new AtomicReference<>();
+
+		aggregator.aggregate(streamingResponse, aggregatedResponseRef::set).blockLast();
+
+		ChatResponse finalResponse = aggregatedResponseRef.get();
+		assertThat(finalResponse).isNotNull();
+
+		AssistantMessage finalAssistantMessage = finalResponse.getResult().getOutput();
+
+		assertThat(finalAssistantMessage).isNotNull();
+		assertThat(finalAssistantMessage.getText()).isEqualTo("Thinking about the weather... ");
+		assertThat(finalAssistantMessage.hasToolCalls()).isTrue();
+		assertThat(finalAssistantMessage.getToolCalls()).hasSize(1);
+
+		ToolCall resultToolCall = finalAssistantMessage.getToolCalls().get(0);
+		assertThat(resultToolCall.id()).isEqualTo("tool-id-123");
+		assertThat(resultToolCall.name()).isEqualTo("getCurrentWeather");
+		assertThat(resultToolCall.arguments()).isEqualTo("{\"location\": \"Seoul\"}");
+	}
+
 }