@@ -4,7 +4,10 @@ import (
44 "context"
55 "errors"
66 "os/exec"
7+ "strings"
78 "testing"
9+
10+ "github.com/entireio/cli/cmd/entire/cli/agent"
811)
912
1013func TestResolveSessionFile (t * testing.T ) {
@@ -26,34 +29,56 @@ func TestProtectedDirs(t *testing.T) {
2629 }
2730}
2831
29- func TestGenerateText_ArrayResponse (t * testing.T ) {
32+ func TestGenerateTextStreaming_SuccessEmitsPhases (t * testing.T ) {
3033 t .Parallel ()
31- ag := & ClaudeCodeAgent {
32- CommandRunner : func (ctx context.Context , _ string , _ ... string ) * exec.Cmd {
33- response := `[{"type":"system","subtype":"init"},{"type":"assistant","message":"Working on it"},{"type":"result","result":"final generated text"}]`
34- return exec .CommandContext (ctx , "sh" , "-c" , "printf '%s' '" + response + "'" )
35- },
36- }
37-
38- result , err := ag .GenerateText (context .Background (), "prompt" , "" )
34+ body := strings .Join ([]string {
35+ `{"type":"system","subtype":"status","status":"requesting"}` ,
36+ `{"type":"stream_event","event":{"type":"message_start","ttft_ms":1500,"message":{"usage":{"input_tokens":10,"cache_read_input_tokens":2000}}}}` ,
37+ `{"type":"stream_event","event":{"type":"content_block_delta","delta":{"type":"text_delta","text":"hello world"}}}` ,
38+ `{"type":"result","subtype":"success","is_error":false,"result":"hello world","duration_ms":1700}` ,
39+ }, "\n " )
40+ ag := newAgentWithStdout (body )
41+ var phases []agent.ProgressPhase
42+ got , err := ag .GenerateTextStreaming (context .Background (), "p" , "" , func (p agent.GenerationProgress ) {
43+ phases = append (phases , p .Phase )
44+ })
3945 if err != nil {
40- t .Fatalf ("unexpected error: %v" , err )
46+ t .Fatalf ("err = %v; want nil" , err )
47+ }
48+ if got != "hello world" {
49+ t .Fatalf ("got = %q; want %q" , got , "hello world" )
50+ }
51+ want := []agent.ProgressPhase {agent .PhaseConnecting , agent .PhaseFirstToken , agent .PhaseGenerating , agent .PhaseDone }
52+ if ! equalPhases (phases , want ) {
53+ t .Fatalf ("phases = %v; want %v" , phases , want )
4154 }
55+ }
4256
43- if result != "final generated text" {
44- t .Fatalf ("GenerateText() = %q, want %q" , result , "final generated text" )
57+ func TestGenerateTextStreaming_EnvelopeErrorReturnsClaudeError (t * testing.T ) {
58+ t .Parallel ()
59+ body := strings .Join ([]string {
60+ `{"type":"system","subtype":"status","status":"requesting"}` ,
61+ `{"type":"result","subtype":"success","is_error":true,"api_error_status":401,"result":"Auth required"}` ,
62+ }, "\n " )
63+ ag := newAgentWithStdout (body )
64+ _ , err := ag .GenerateTextStreaming (context .Background (), "p" , "" , nil )
65+ var ce * ClaudeError
66+ if ! errors .As (err , & ce ) {
67+ t .Fatalf ("err = %v; want *ClaudeError" , err )
68+ }
69+ if ce .Kind != ClaudeErrorAuth {
70+ t .Fatalf ("Kind = %v; want %v" , ce .Kind , ClaudeErrorAuth )
4571 }
4672}
4773
48- func TestGenerateText_EnvelopeErrorReturnsClaudeError (t * testing.T ) {
74+ func TestGenerateTextStreaming_StderrAuthFallback (t * testing.T ) {
4975 t .Parallel ()
5076 ag := & ClaudeCodeAgent {
5177 CommandRunner : func (ctx context.Context , _ string , _ ... string ) * exec.Cmd {
52- response := `{"type":"result","subtype":"success","is_error":true,"api_error_status":401,"result":"Auth required"}`
53- return exec .CommandContext (ctx , "sh" , "-c" , "printf '%s' '" + response + "'" )
78+ return exec .CommandContext (ctx , "sh" , "-c" , "printf 'Invalid API key' 1>&2; exit 2" )
5479 },
5580 }
56- _ , err := ag .GenerateText (context .Background (), "prompt " , "" )
81+ _ , err := ag .GenerateTextStreaming (context .Background (), "p " , "" , nil )
5782 var ce * ClaudeError
5883 if ! errors .As (err , & ce ) {
5984 t .Fatalf ("err = %v; want *ClaudeError" , err )
@@ -63,14 +88,14 @@ func TestGenerateText_EnvelopeErrorReturnsClaudeError(t *testing.T) {
6388 }
6489}
6590
66- func TestGenerateText_CLIMissing (t * testing.T ) {
91+ func TestGenerateTextStreaming_CLIMissing (t * testing.T ) {
6792 t .Parallel ()
6893 ag := & ClaudeCodeAgent {
6994 CommandRunner : func (ctx context.Context , _ string , _ ... string ) * exec.Cmd {
7095 return exec .CommandContext (ctx , "/nonexistent/binary/that/does/not/exist" )
7196 },
7297 }
73- _ , err := ag .GenerateText (context .Background (), "prompt " , "" )
98+ _ , err := ag .GenerateTextStreaming (context .Background (), "p " , "" , nil )
7499 var ce * ClaudeError
75100 if ! errors .As (err , & ce ) {
76101 t .Fatalf ("err = %v; want *ClaudeError" , err )
@@ -80,19 +105,37 @@ func TestGenerateText_CLIMissing(t *testing.T) {
80105 }
81106}
82107
83- func TestGenerateText_StderrAuthFallback (t * testing.T ) {
108+ func TestGenerateText_DelegatesToStreaming (t * testing.T ) {
84109 t .Parallel ()
85- ag := & ClaudeCodeAgent {
110+ body := `{"type":"result","subtype":"success","is_error":false,"result":"ok"}`
111+ ag := newAgentWithStdout (body )
112+ got , err := ag .GenerateText (context .Background (), "p" , "" )
113+ if err != nil {
114+ t .Fatalf ("err = %v; want nil" , err )
115+ }
116+ if got != "ok" {
117+ t .Fatalf ("got = %q; want %q" , got , "ok" )
118+ }
119+ }
120+
121+ // newAgentWithStdout returns a ClaudeCodeAgent whose CommandRunner produces a
122+ // subprocess that prints the given body to stdout and exits 0.
123+ func newAgentWithStdout (body string ) * ClaudeCodeAgent {
124+ return & ClaudeCodeAgent {
86125 CommandRunner : func (ctx context.Context , _ string , _ ... string ) * exec.Cmd {
87- return exec .CommandContext (ctx , "sh" , "-c" , "printf 'Invalid API key' 1>&2; exit 2 " )
126+ return exec .CommandContext (ctx , "sh" , "-c" , "cat <<'ENDOFSTREAM' \n " + body + " \n ENDOFSTREAM " )
88127 },
89128 }
90- _ , err := ag .GenerateText (context .Background (), "prompt" , "" )
91- var ce * ClaudeError
92- if ! errors .As (err , & ce ) {
93- t .Fatalf ("err = %v; want *ClaudeError" , err )
129+ }
130+
131+ func equalPhases (a , b []agent.ProgressPhase ) bool {
132+ if len (a ) != len (b ) {
133+ return false
94134 }
95- if ce .Kind != ClaudeErrorAuth {
96- t .Fatalf ("Kind = %v; want %v" , ce .Kind , ClaudeErrorAuth )
135+ for i := range a {
136+ if a [i ] != b [i ] {
137+ return false
138+ }
97139 }
140+ return true
98141}
0 commit comments