mirror of
https://github.com/coder/coder.git
synced 2026-06-02 20:48:20 +00:00
Paweł Banaszewski
e00e85765b
This PR merges code from `coder/aibridge` repository into `coder/coder`. It was split into 4 PRs for easier review but stacked PRs will need to be merged into this PR so all checks pass. * https://github.com/coder/coder/pull/24190 -> raw code copy (this PR, before merging PRs on top of it, it was just 1 commit: https://github.com/coder/coder/commit/70d33f33200c7e77df910957595715f81f9bec24) * https://github.com/coder/coder/pull/24570 -> update imports in `coder/coder` to use copied code * https://github.com/coder/coder/pull/24586 -> linter fixes and CI integration (also added README.md) * https://github.com/coder/coder/pull/24571 -> added exclude to scripts/check_emdash.sh check Original PR message (before PR squash): Moves coder/aibridge code into coder/coder repository. Omitted files: - `go.mod`, `go.sum`, `.gitignore`, `.github/workflows/ci.yml,` `Makefile`, `LICENSE`, `README.md` (modified README.md is added later) - `.github`, `example`, `buildinfo,` `scripts` directories Simple verification script (will list omitted files) ``` tmp=$(mktemp -d) echo "$tmp" git clone --depth=1 https://github.com/coder/aibridge "$tmp/aibridge" git clone --depth=1 --branch pb/aibridge-code-move https://github.com/coder/coder "$tmp/coder" diff -rq --exclude=.git "$tmp/aibridge" "$tmp/coder/aibridge" # rm -rf "$tmp" ```
111 lines
2.8 KiB
Go
111 lines
2.8 KiB
Go
package eventstream_test
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"net"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"cdr.dev/slog/v3"
|
|
"cdr.dev/slog/v3/sloggers/sloghuman"
|
|
"cdr.dev/slog/v3/sloggers/slogtest"
|
|
"github.com/coder/coder/v2/aibridge/intercept/eventstream"
|
|
"github.com/coder/quartz"
|
|
)
|
|
|
|
// clockAdvancingFlusher wraps httptest.ResponseRecorder and advances the mock
|
|
// clock on each Flush call, simulating a slow client without real sleeping.
|
|
type clockAdvancingFlusher struct {
|
|
*httptest.ResponseRecorder
|
|
clk *quartz.Mock
|
|
advance time.Duration
|
|
}
|
|
|
|
func (f *clockAdvancingFlusher) Flush() {
|
|
f.clk.Advance(f.advance)
|
|
f.ResponseRecorder.Flush()
|
|
}
|
|
|
|
// Hijack satisfies the FullResponseWriter lint rule.
|
|
func (*clockAdvancingFlusher) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
|
return nil, nil, nil
|
|
}
|
|
|
|
func TestEventStream_LogsWarning_WhenFlushIsSlow(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
var buf strings.Builder
|
|
logger := slogtest.Make(t, nil).AppendSinks(sloghuman.Sink(&buf)).Leveled(slog.LevelWarn)
|
|
ctx := context.Background()
|
|
clk := quartz.NewMock(t)
|
|
|
|
stream := eventstream.NewEventStream(ctx, logger, nil, clk)
|
|
|
|
w := &clockAdvancingFlusher{
|
|
ResponseRecorder: httptest.NewRecorder(),
|
|
clk: clk,
|
|
advance: eventstream.SlowFlushThreshold + time.Millisecond,
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "/", nil)
|
|
require.NoError(t, err)
|
|
req.RemoteAddr = "192.0.2.1:12345"
|
|
req.Header.Set("User-Agent", "test-agent/1.0")
|
|
|
|
done := make(chan struct{})
|
|
go func() {
|
|
defer close(done)
|
|
stream.Start(w, req)
|
|
}()
|
|
|
|
stream.InitiateStream(w)
|
|
require.NoError(t, stream.SendRaw(ctx, []byte("data: hello\n\n")))
|
|
require.NoError(t, stream.Shutdown(ctx))
|
|
<-done
|
|
|
|
require.Contains(t, buf.String(), "slow client detected")
|
|
require.Contains(t, buf.String(), "192.0.2.1")
|
|
require.Contains(t, buf.String(), "test-agent/1.0")
|
|
require.Contains(t, buf.String(), "payload_size=13")
|
|
}
|
|
|
|
func TestEventStream_NoWarning_WhenFlushIsFast(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
var buf strings.Builder
|
|
logger := slogtest.Make(t, nil).AppendSinks(sloghuman.Sink(&buf)).Leveled(slog.LevelWarn)
|
|
ctx := context.Background()
|
|
clk := quartz.NewMock(t)
|
|
|
|
stream := eventstream.NewEventStream(ctx, logger, nil, clk)
|
|
|
|
// No clock advance, flush duration stays at 0, below threshold.
|
|
w := &clockAdvancingFlusher{
|
|
ResponseRecorder: httptest.NewRecorder(),
|
|
clk: clk,
|
|
advance: 0,
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "/", nil)
|
|
require.NoError(t, err)
|
|
|
|
done := make(chan struct{})
|
|
go func() {
|
|
defer close(done)
|
|
stream.Start(w, req)
|
|
}()
|
|
|
|
stream.InitiateStream(w)
|
|
require.NoError(t, stream.SendRaw(ctx, []byte("data: hello\n\n")))
|
|
require.NoError(t, stream.Shutdown(ctx))
|
|
<-done
|
|
|
|
require.Empty(t, buf.String())
|
|
}
|