Repository: google/adk-go
File: agent/remoteagent/a2a_agent_test.go (Go)

Code before:
	card := &a2a.AgentCard{PreferredTransport: a2a.TransportProtocolJSONRPC, URL: server.URL, Capabilities: a2a.AgentCapabilities{Streaming: true}}
	agent, err := NewA2A(A2AConfig{Name: name, AgentCard: card})
	if err != nil {
		t.Fatalf("remoteagent.NewA2A() error = %v", err)
	}
	return agent
}

func newInvocationContext(t *testing.T, events []*session.Event) agent.InvocationContext {
	t.Helper()
	ctx := t.Context()
	service := session.InMemoryService()
	resp, err := service.Create(ctx, &session.CreateRequest{AppName: t.Name(), UserID: "test"})
	if err != nil {
		t.Fatalf("sessionService.Create() error = %v", err)
	}
	for _, event := range events {
		if err := service.AppendEvent(ctx, resp.Session, event); err != nil {
			t.Fatalf("sessionService.AppendEvent() error = %v", err)
		}
	}
	ic := icontext.NewInvocationContext(ctx, icontext.InvocationContextParams{Session: resp.Session})
	return ic
}

func newInvocationContextWithStreamingMode(t *testing.T, events []*session.Event, noStreaming bool) agent.InvocationContext {
	t.Helper()
	ctx := t.Context()
	service := session.InMemoryService()
	resp, err := service.Create(ctx, &session.CreateRequest{AppName: t.Name(), UserID: "test"})
	if err != nil {
		t.Fatalf("sessionService.Create() error = %v", err)
	}
	for _, event := range events {
		if err := service.AppendEvent(ctx, resp.Session, event); err != nil {
			t.Fatalf("sessionService.AppendEvent() error = %v", err)
		}
	}

	streamingMode := agent.StreamingModeSSE
	if noStreaming {
		streamingMode = agent.StreamingModeNone
	}

	ic := icontext.NewInvocationContext(ctx, icontext.InvocationContextParams{
		Session: resp.Session,
		RunConfig: &agent.RunConfig{
			StreamingMode: streamingMode,
		},
	})
	return ic

Review comment (refactor):
I'd rather reduce context setup code duplication:

<details>

```diff
diff --git a/agent/remoteagent/a2a_agent_test.go b/agent/remoteagent/a2a_agent_test.go
index 9dbda44..96a2263 100644
--- a/agent/remoteagent/a2a_agent_test.go
+++ b/agent/remoteagent/a2a_agent_test.go
@@ -73,23 +73,10 @@ func newA2ARemoteAgent(t *testing.T, name string, server *httptest.Server) agent
 }

 func newInvocationContext(t *testing.T, events []*session.Event) agent.InvocationContext {
-	t.Helper()
-	ctx := t.Context()
-	service := session.InMemoryService()
-	resp, err := service.Create(ctx, &session.CreateRequest{AppName: t.Name(), UserID: "test"})
-	if err != nil {
-		t.Fatalf("sessionService.Create() error = %v", err)
-	}
-	for _, event := range events {
-		if err := service.AppendEvent(ctx, resp.Session, event); err != nil {
-			t.Fatalf("sessionService.AppendEvent() error = %v", err)
-		}
-	}
-	ic := icontext.NewInvocationContext(ctx, icontext.InvocationContextParams{Session: resp.Session})
-	return ic
+	return newInvocationContextWithStreamingMode(t, events, agent.StreamingModeSSE)
 }

-func newInvocationContextWithStreamingMode(t *testing.T, events []*session.Event, noStreaming bool) agent.InvocationContext {
+func newInvocationContextWithStreamingMode(t *testing.T, events []*session.Event, streamingMode agent.StreamingMode) agent.InvocationContext {
 	t.Helper()
 	ctx := t.Context()
 	service := session.InMemoryService()
@@ -103,11 +90,6 @@ func newInvocationContextWithStreamingMode(t *testing.T, events []*session.Event
 		}
 	}

-	streamingMode := agent.StreamingModeSSE
-	if noStreaming {
-		streamingMode = agent.StreamingModeNone
-	}
-
 	ic := icontext.NewInvocationContext(ctx, icontext.InvocationContextParams{
 		Session: resp.Session,
 		RunConfig: &agent.RunConfig{
@@ -383,7 +365,11 @@ func TestRemoteAgent_ADK2ADK(t *testing.T) {
 			executor := newADKEventReplay(t, tc.remoteEvents)
 			remoteAgent := newA2ARemoteAgent(t, "a2a", startA2AServer(executor))

-			ictx := newInvocationContextWithStreamingMode(t, []*session.Event{newUserHello()}, tc.noStreaming)
+			mode := agent.StreamingModeSSE
+			if tc.noStreaming {
+				mode = agent.StreamingModeNone
+			}
+			ictx := newInvocationContextWithStreamingMode(t, []*session.Event{newUserHello()}, mode)
 			gotEvents, err := runAndCollect(ictx, remoteAgent)
 			if err != nil {
 				t.Fatalf("agent.Run() error = %v", err)
```
</details>

Code after:
	card := &a2a.AgentCard{PreferredTransport: a2a.TransportProtocolJSONRPC, URL: server.URL, Capabilities: a2a.AgentCapabilities{Streaming: true}}
	agent, err := NewA2A(A2AConfig{Name: name, AgentCard: card})
	if err != nil {
		t.Fatalf("remoteagent.NewA2A() error = %v", err)
	}
	return agent
}

func newInvocationContext(t *testing.T, events []*session.Event) agent.InvocationContext {
	return newInvocationContextWithStreamingMode(t, events, agent.StreamingModeSSE)
}

func newInvocationContextWithStreamingMode(t *testing.T, events []*session.Event, streamingMode agent.StreamingMode) agent.InvocationContext {
	t.Helper()
	ctx := t.Context()
	service := session.InMemoryService()
	resp, err := service.Create(ctx, &session.CreateRequest{AppName: t.Name(), UserID: "test"})
	if err != nil {
		t.Fatalf("sessionService.Create() error = %v", err)
	}
	for _, event := range events {
		if err := service.AppendEvent(ctx, resp.Session, event); err != nil {
			t.Fatalf("sessionService.AppendEvent() error = %v", err)
		}
	}

	ic := icontext.NewInvocationContext(ctx, icontext.InvocationContextParams{
		Session: resp.Session,
		RunConfig: &agent.RunConfig{
			StreamingMode: streamingMode,
		},
	})
	return ic
}

func runAndCollect(ic agent.InvocationContext, agnt agent.Agent) ([]*session.Event, error) {
	var collected []*session.Event
	for ev, err := range agnt.Run(ic) {
		if err != nil {
			return collected, err
		}
		collected = append(collected, ev)
	}
	return collected, nil
}

func toLLMResponses(events []*session.Event) []model.LLMResponse {
	var result []model.LLMResponse
	for _, v := range events {
		result = append(result, v.LLMResponse)
	}