95 lines
2.7 KiB
Go
95 lines
2.7 KiB
Go
package logging
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
type roundTripFunc func(*http.Request) (*http.Response, error)
|
|
|
|
func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
|
|
|
|
type upgradeStream struct{ bytes.Buffer }
|
|
|
|
func (*upgradeStream) Close() error { return nil }
|
|
|
|
func TestHTTPUpgradeRetainsOriginalReadWriteStream(t *testing.T) {
|
|
installTestLogger(t, true)
|
|
stream := &upgradeStream{}
|
|
client := HTTPClient(&http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
|
return &http.Response{StatusCode: 101, Header: make(http.Header), Body: stream}, nil
|
|
})}, "agent", "asr")
|
|
req, _ := http.NewRequest(http.MethodGet, "http://example.test/ws", nil)
|
|
res, err := client.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if res.Body != stream {
|
|
t.Fatal("WebSocket upgraded stream was decorated")
|
|
}
|
|
if _, ok := res.Body.(io.ReadWriteCloser); !ok {
|
|
t.Fatal("WebSocket writer lost")
|
|
}
|
|
}
|
|
|
|
func TestDefaultTransportObservationClonesBaseAndDoesNotDoubleLog(t *testing.T) {
|
|
previous := http.DefaultTransport
|
|
t.Cleanup(func() { http.DefaultTransport = previous })
|
|
http.DefaultTransport = &http.Transport{}
|
|
ObserveDefaultHTTP("agent")
|
|
clone, err := CloneDefaultTransport()
|
|
if err != nil || clone == nil {
|
|
t.Fatal(err)
|
|
}
|
|
out := installTestLogger(t, false)
|
|
base := roundTripFunc(func(*http.Request) (*http.Response, error) {
|
|
return &http.Response{StatusCode: 200, Header: make(http.Header), Body: http.NoBody}, nil
|
|
})
|
|
http.DefaultTransport = &httpObserver{next: base, local: "agent"}
|
|
client := HTTPClient(nil, "agent", "llm")
|
|
req, _ := http.NewRequest(http.MethodGet, "http://example.test/", nil)
|
|
res, err := client.Do(req)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
res.Body.Close()
|
|
if strings.Count(out.String(), `"stage":"request_start"`) != 1 {
|
|
t.Fatal(out.String())
|
|
}
|
|
if _, err = CloneDefaultTransport(); err == nil {
|
|
t.Fatal("custom unsupported transport must be explicit")
|
|
}
|
|
}
|
|
|
|
func TestHTTPReadFailureIsReportedWithoutLeakingIncompleteBody(t *testing.T) {
|
|
out := installTestLogger(t, true)
|
|
expected := errors.New("read failed")
|
|
body := newObservedBody(context.Background(), &brokenBody{err: expected}, "application/json", "HTTP response data")
|
|
if _, err := io.ReadAll(body); !errors.Is(err, expected) {
|
|
t.Fatal(err)
|
|
}
|
|
body.Close()
|
|
if strings.Contains(out.String(), "private") || !strings.Contains(out.String(), "body_read_incomplete") || !strings.Contains(out.String(), "read failed") {
|
|
t.Fatal(out.String())
|
|
}
|
|
}
|
|
|
|
type brokenBody struct {
|
|
err error
|
|
read bool
|
|
}
|
|
|
|
func (b *brokenBody) Read(p []byte) (int, error) {
|
|
if b.read {
|
|
return 0, b.err
|
|
}
|
|
b.read = true
|
|
return copy(p, `{"api_key":"private`), b.err
|
|
}
|
|
func (*brokenBody) Close() error { return nil }
|