|
26 | 26 | import static com.google.common.truth.Truth.assertThat; |
27 | 27 | import static java.nio.charset.StandardCharsets.UTF_8; |
28 | 28 | import static java.util.Arrays.stream; |
| 29 | +import static java.util.concurrent.TimeUnit.MILLISECONDS; |
29 | 30 | import static org.mockito.ArgumentMatchers.any; |
30 | 31 | import static org.mockito.Mockito.CALLS_REAL_METHODS; |
31 | 32 | import static org.mockito.Mockito.mock; |
32 | 33 | import static org.mockito.Mockito.never; |
33 | 34 | import static org.mockito.Mockito.verify; |
34 | 35 | import static org.mockito.Mockito.when; |
35 | 36 |
|
| 37 | +import com.google.adk.agents.BaseAgent; |
36 | 38 | import com.google.adk.agents.InvocationContext; |
37 | 39 | import com.google.adk.agents.LiveRequestQueue; |
38 | 40 | import com.google.adk.agents.LlmAgent; |
|
43 | 45 | import com.google.adk.flows.llmflows.Functions; |
44 | 46 | import com.google.adk.models.LlmResponse; |
45 | 47 | import com.google.adk.plugins.BasePlugin; |
| 48 | +import com.google.adk.sessions.BaseSessionService; |
46 | 49 | import com.google.adk.sessions.Session; |
47 | 50 | import com.google.adk.sessions.SessionKey; |
48 | 51 | import com.google.adk.summarizer.EventsCompactionConfig; |
@@ -851,6 +854,45 @@ public void beforeRunCallback_withStateDelta_seesMergedState() { |
851 | 854 | assertThat(sessionInCallback.state()).containsEntry("number", 123); |
852 | 855 | } |
853 | 856 |
|
| 857 | + @Test |
| 858 | + public void runAsync_ensureEventsAreAppendedInOrder() throws Exception { |
| 859 | + Event event1 = TestUtils.createEvent("1"); |
| 860 | + Event event2 = TestUtils.createEvent("2"); |
| 861 | + BaseAgent mockAgent = TestUtils.createSubAgent("test agent", event1, event2); |
| 862 | + |
| 863 | + BaseSessionService mockSessionService = mock(BaseSessionService.class); |
| 864 | + |
| 865 | + when(mockSessionService.getSession(any(), any(), any(), any())).thenReturn(Maybe.just(session)); |
| 866 | + when(mockSessionService.appendEvent(any(), any())) |
| 867 | + .thenAnswer( |
| 868 | + invocation -> { |
| 869 | + Event eventArg = invocation.getArgument(1); |
| 870 | + Single<Event> result = Single.just(eventArg); |
| 871 | + if (eventArg.id().equals("1")) { |
| 872 | + // Artificially delay the first event to ensure it is appended first. |
| 873 | + return result.delay(100, MILLISECONDS); |
| 874 | + } |
| 875 | + return result; |
| 876 | + }); |
| 877 | + |
| 878 | + Runner mockRunner = |
| 879 | + Runner.builder() |
| 880 | + .agent(mockAgent) |
| 881 | + .appName("test") |
| 882 | + .sessionService(mockSessionService) |
| 883 | + .build(); |
| 884 | + |
| 885 | + List<Event> results = |
| 886 | + mockRunner |
| 887 | + .runAsync("user", session.id(), createContent("user message")) |
| 888 | + .toList() |
| 889 | + .blockingGet(); |
| 890 | + |
| 891 | + assertThat(simplifyEvents(results)) |
| 892 | + .containsExactly("author: content for event 1", "author: content for event 2") |
| 893 | + .inOrder(); |
| 894 | + } |
| 895 | + |
854 | 896 | private Content createContent(String text) { |
855 | 897 | return Content.builder().parts(Part.builder().text(text).build()).build(); |
856 | 898 | } |
|
0 commit comments