first commit

pigo/internal/agentcore/event_stream_test.go

@@ -0,0 +1,121 @@
+package agentcore
+
+import (
+	"context"
+	"testing"
+	"time"
+)
+
+// TestEventStreamNormalCompletion drives a producer that emits events and sets
+// a result, then verifies the consumer sees every event and Result yields the
+// captured value.
+func TestEventStreamNormalCompletion(t *testing.T) {
+	s := NewEventStream[AgentEvent, []AgentMessage](0)
+	want := []AgentMessage{
+		UserMessage{RoleField: RoleUser, Content: ContentList{NewTextContent("hi")}},
+	}
+
+	go func() {
+		ctx := context.Background()
+		_ = s.Emit(ctx, TurnStartEvent{})
+		_ = s.Emit(ctx, AgentEndEvent{Messages: want})
+		s.SetResult(want)
+		s.Close()
+	}()
+
+	var got int
+	for range s.Events() {
+		got++
+	}
+	if got != 2 {
+		t.Fatalf("want 2 events, got %d", got)
+	}
+	res, err := s.Result(context.Background())
+	if err != nil {
+		t.Fatalf("result: %v", err)
+	}
+	if len(res) != 1 || res[0].Role() != RoleUser {
+		t.Fatalf("result payload wrong: %+v", res)
+	}
+}
+
+// TestEventStreamIsCompleteCallback verifies the isComplete/extractResult
+// callbacks auto-capture the result on the terminal event.
+func TestEventStreamIsCompleteCallback(t *testing.T) {
+	s := NewEventStream[AgentEvent, []AgentMessage](4)
+	s.IsComplete = func(e AgentEvent) bool { return e.EventType() == EventAgentEnd }
+	s.ExtractResult = func(e AgentEvent) []AgentMessage { return e.(AgentEndEvent).Messages }
+
+	msgs := []AgentMessage{AssistantMessage{RoleField: RoleAssistant}}
+	go func() {
+		ctx := context.Background()
+		_ = s.Emit(ctx, MessageStartEvent{})
+		_ = s.Emit(ctx, AgentEndEvent{Messages: msgs})
+		s.Close()
+	}()
+
+	for range s.Events() {
+	}
+	res, err := s.Result(context.Background())
+	if err != nil {
+		t.Fatalf("result: %v", err)
+	}
+	if len(res) != 1 {
+		t.Fatalf("want 1 msg from extractResult, got %d", len(res))
+	}
+}
+
+// TestEventStreamCancellation verifies that a cancelled context unblocks a
+// producer stuck on Emit (consumer stopped reading) and that Result returns the
+// context error.
+func TestEventStreamCancellation(t *testing.T) {
+	s := NewEventStream[AgentEvent, []AgentMessage](0)
+	ctx, cancel := context.WithCancel(context.Background())
+
+	emitErr := make(chan error, 1)
+	go func() {
+		// First emit has no consumer; it blocks until cancel.
+		emitErr <- s.Emit(ctx, TurnStartEvent{})
+	}()
+
+	// Give the producer a moment to block on the send, then cancel.
+	time.Sleep(20 * time.Millisecond)
+	cancel()
+
+	select {
+	case err := <-emitErr:
+		if err == nil {
+			t.Fatal("expected Emit to return ctx error on cancellation")
+		}
+	case <-time.After(time.Second):
+		t.Fatal("Emit did not unblock after cancel (goroutine leak)")
+	}
+
+	// Result with a cancelled context returns promptly with the ctx error.
+	if _, err := s.Result(ctx); err == nil {
+		t.Fatal("expected Result to return ctx error")
+	}
+}
+
+// TestEventStreamIncompleteClose verifies Close without a result yields
+// ErrStreamIncomplete.
+func TestEventStreamIncompleteClose(t *testing.T) {
+	s := NewEventStream[AgentEvent, []AgentMessage](1)
+	s.Close()
+	if _, err := s.Result(context.Background()); err != ErrStreamIncomplete {
+		t.Fatalf("want ErrStreamIncomplete, got %v", err)
+	}
+}
+
+// TestEventStreamSetErrorWins verifies SetError is reported and later SetResult
+// is ignored (first outcome wins).
+func TestEventStreamSetErrorWins(t *testing.T) {
+	s := NewEventStream[AgentEvent, []AgentMessage](1)
+	sentinel := context.Canceled
+	s.SetError(sentinel)
+	s.SetResult(nil)
+	s.Close()
+	if _, err := s.Result(context.Background()); err != sentinel {
+		t.Fatalf("want sentinel error, got %v", err)
+	}
+}