Skip to content

Commit 2509bd0

Browse files
ibetitsmikeibetitsmike
committed
testing and fixes
1 parent 192b0e2 commit 2509bd0

File tree

8 files changed

+1610
-10
lines changed

8 files changed

+1610
-10
lines changed

‎agent/immortalstream.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -54,13 +54,14 @@ func (h *immortalStreamHandler) createStream(w http.ResponseWriter, r *http.Requ
5454

5555
resp, err := h.manager.CreateStream(ctx, req)
5656
if err != nil {
57-
if xerrors.Is(err, xerrors.Errorf("Too many Immortal Streams")) {
57+
errMsg := err.Error()
58+
if errMsg == "Too many Immortal Streams" {
5859
httpapi.Write(ctx, w, http.StatusServiceUnavailable, codersdk.Response{
5960
Message: "Too many Immortal Streams",
6061
})
6162
return
6263
}
63-
if xerrors.Is(err, xerrors.Errorf("The connection was refused")) {
64+
if errMsg == "The connection was refused" {
6465
httpapi.Write(ctx, w, http.StatusNotFound, codersdk.Response{
6566
Message: "The connection was refused",
6667
})
@@ -182,7 +183,7 @@ func (h *immortalStreamHandler) deleteStream(w http.ResponseWriter, r *http.Requ
182183

183184
err = h.manager.DeleteStream(streamID)
184185
if err != nil {
185-
if xerrors.Is(err, xerrors.Errorf("stream not found")) {
186+
if err.Error() == "stream not found" {
186187
httpapi.Write(ctx, w, http.StatusNotFound, codersdk.Response{
187188
Message: "Stream not found",
188189
})

‎agent/immortalstream_test.go

Lines changed: 366 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,366 @@
1+
package agent
2+
3+
import (
4+
"bytes"
5+
"context"
6+
"encoding/json"
7+
"io"
8+
"net"
9+
"net/http"
10+
"net/http/httptest"
11+
"strings"
12+
"testing"
13+
"time"
14+
15+
"github.com/go-chi/chi/v5"
16+
"github.com/google/uuid"
17+
"golang.org/x/xerrors"
18+
19+
"cdr.dev/slog/sloggers/slogtest"
20+
"github.com/coder/coder/v2/immortalstream"
21+
)
22+
23+
func TestImmortalStreamHandler(t *testing.T) {
24+
t.Parallel()
25+
26+
t.Run("CreateStream", func(t *testing.T) {
27+
t.Parallel()
28+
29+
logger := slogtest.Make(t, nil)
30+
dialer := func(network, address string) (net.Conn, error) {
31+
return &mockConn{}, nil
32+
}
33+
manager := immortalstream.NewManager(logger, dialer)
34+
defer manager.Close()
35+
36+
handler := newImmortalStreamHandler(manager)
37+
38+
req := immortalstream.CreateStreamRequest{TCPPort: 8080}
39+
body, _ := json.Marshal(req)
40+
41+
httpReq := httptest.NewRequest("POST", "/api/v0/immortal-stream", bytes.NewReader(body))
42+
httpReq.Header.Set("Content-Type", "application/json")
43+
44+
rr := httptest.NewRecorder()
45+
handler.createStream(rr, httpReq)
46+
47+
if rr.Code != http.StatusCreated {
48+
t.Fatalf("expected status 201, got %d", rr.Code)
49+
}
50+
51+
var resp immortalstream.CreateStreamResponse
52+
err := json.Unmarshal(rr.Body.Bytes(), &resp)
53+
if err != nil {
54+
t.Fatalf("failed to unmarshal response: %v", err)
55+
}
56+
57+
if resp.ID == uuid.Nil {
58+
t.Fatalf("expected valid stream ID")
59+
}
60+
if resp.Name == "" {
61+
t.Fatalf("expected non-empty stream name")
62+
}
63+
})
64+
65+
t.Run("CreateStreamInvalidJSON", func(t *testing.T) {
66+
t.Parallel()
67+
68+
logger := slogtest.Make(t, nil)
69+
manager := immortalstream.NewManager(logger, nil)
70+
defer manager.Close()
71+
72+
handler := newImmortalStreamHandler(manager)
73+
74+
httpReq := httptest.NewRequest("POST", "/api/v0/immortal-stream", strings.NewReader("invalid json"))
75+
httpReq.Header.Set("Content-Type", "application/json")
76+
77+
rr := httptest.NewRecorder()
78+
handler.createStream(rr, httpReq)
79+
80+
if rr.Code != http.StatusBadRequest {
81+
t.Fatalf("expected status 400, got %d", rr.Code)
82+
}
83+
})
84+
85+
t.Run("CreateStreamConnectionRefused", func(t *testing.T) {
86+
t.Parallel()
87+
88+
logger := slogtest.Make(t, nil)
89+
dialer := func(network, address string) (net.Conn, error) {
90+
return nil, xerrors.New("connection refused")
91+
}
92+
manager := immortalstream.NewManager(logger, dialer)
93+
defer manager.Close()
94+
95+
handler := newImmortalStreamHandler(manager)
96+
97+
req := immortalstream.CreateStreamRequest{TCPPort: 8080}
98+
body, _ := json.Marshal(req)
99+
100+
httpReq := httptest.NewRequest("POST", "/api/v0/immortal-stream", bytes.NewReader(body))
101+
httpReq.Header.Set("Content-Type", "application/json")
102+
103+
rr := httptest.NewRecorder()
104+
handler.createStream(rr, httpReq)
105+
106+
if rr.Code != http.StatusNotFound {
107+
t.Fatalf("expected status 404, got %d", rr.Code)
108+
}
109+
})
110+
111+
t.Run("ListStreams", func(t *testing.T) {
112+
t.Parallel()
113+
114+
logger := slogtest.Make(t, nil)
115+
dialer := func(network, address string) (net.Conn, error) {
116+
return &mockConn{}, nil
117+
}
118+
manager := immortalstream.NewManager(logger, dialer)
119+
defer manager.Close()
120+
121+
// Create a test stream
122+
createReq := immortalstream.CreateStreamRequest{TCPPort: 8080}
123+
_, err := manager.CreateStream(context.Background(), createReq)
124+
if err != nil {
125+
t.Fatalf("failed to create test stream: %v", err)
126+
}
127+
128+
handler := newImmortalStreamHandler(manager)
129+
130+
httpReq := httptest.NewRequest("GET", "/api/v0/immortal-stream", nil)
131+
rr := httptest.NewRecorder()
132+
handler.listStreams(rr, httpReq)
133+
134+
if rr.Code != http.StatusOK {
135+
t.Fatalf("expected status 200, got %d", rr.Code)
136+
}
137+
138+
var streams []*immortalstream.Stream
139+
err = json.Unmarshal(rr.Body.Bytes(), &streams)
140+
if err != nil {
141+
t.Fatalf("failed to unmarshal response: %v", err)
142+
}
143+
144+
if len(streams) != 1 {
145+
t.Fatalf("expected 1 stream, got %d", len(streams))
146+
}
147+
if streams[0].TCPPort != 8080 {
148+
t.Fatalf("expected port 8080, got %d", streams[0].TCPPort)
149+
}
150+
})
151+
152+
t.Run("DeleteStream", func(t *testing.T) {
153+
t.Parallel()
154+
155+
logger := slogtest.Make(t, nil)
156+
dialer := func(network, address string) (net.Conn, error) {
157+
return &mockConn{}, nil
158+
}
159+
manager := immortalstream.NewManager(logger, dialer)
160+
defer manager.Close()
161+
162+
// Create a test stream
163+
createReq := immortalstream.CreateStreamRequest{TCPPort: 8080}
164+
createResp, err := manager.CreateStream(context.Background(), createReq)
165+
if err != nil {
166+
t.Fatalf("failed to create test stream: %v", err)
167+
}
168+
169+
handler := newImmortalStreamHandler(manager)
170+
171+
// Create chi router to handle URL parameters
172+
r := chi.NewRouter()
173+
r.Delete("/{id}", handler.deleteStream)
174+
175+
httpReq := httptest.NewRequest("DELETE", "/"+createResp.ID.String(), nil)
176+
rr := httptest.NewRecorder()
177+
r.ServeHTTP(rr, httpReq)
178+
179+
if rr.Code != http.StatusNoContent {
180+
t.Fatalf("expected status 204, got %d", rr.Code)
181+
}
182+
183+
// Verify stream was deleted
184+
streams := manager.ListStreams()
185+
if len(streams) != 0 {
186+
t.Fatalf("expected 0 streams after deletion, got %d", len(streams))
187+
}
188+
})
189+
190+
t.Run("DeleteStreamNotFound", func(t *testing.T) {
191+
t.Parallel()
192+
193+
logger := slogtest.Make(t, nil)
194+
manager := immortalstream.NewManager(logger, nil)
195+
defer manager.Close()
196+
197+
handler := newImmortalStreamHandler(manager)
198+
199+
// Create chi router to handle URL parameters
200+
r := chi.NewRouter()
201+
r.Delete("/{id}", handler.deleteStream)
202+
203+
nonExistentID := uuid.New()
204+
httpReq := httptest.NewRequest("DELETE", "/"+nonExistentID.String(), nil)
205+
rr := httptest.NewRecorder()
206+
r.ServeHTTP(rr, httpReq)
207+
208+
if rr.Code != http.StatusNotFound {
209+
t.Fatalf("expected status 404, got %d", rr.Code)
210+
}
211+
})
212+
213+
t.Run("ConnectToStreamMissingUpgradeHeaders", func(t *testing.T) {
214+
t.Parallel()
215+
216+
logger := slogtest.Make(t, nil)
217+
dialer := func(network, address string) (net.Conn, error) {
218+
return &mockConn{}, nil
219+
}
220+
manager := immortalstream.NewManager(logger, dialer)
221+
defer manager.Close()
222+
223+
// Create a test stream
224+
createReq := immortalstream.CreateStreamRequest{TCPPort: 8080}
225+
createResp, err := manager.CreateStream(context.Background(), createReq)
226+
if err != nil {
227+
t.Fatalf("failed to create test stream: %v", err)
228+
}
229+
230+
handler := newImmortalStreamHandler(manager)
231+
232+
// Create chi router to handle URL parameters
233+
r := chi.NewRouter()
234+
r.Get("/{id}", handler.connectToStream)
235+
236+
httpReq := httptest.NewRequest("GET", "/"+createResp.ID.String(), nil)
237+
// Missing upgrade headers
238+
rr := httptest.NewRecorder()
239+
r.ServeHTTP(rr, httpReq)
240+
241+
if rr.Code != http.StatusBadRequest {
242+
t.Fatalf("expected status 400, got %d", rr.Code)
243+
}
244+
})
245+
246+
t.Run("ConnectToStreamInvalidSequenceNumber", func(t *testing.T) {
247+
t.Parallel()
248+
249+
logger := slogtest.Make(t, nil)
250+
dialer := func(network, address string) (net.Conn, error) {
251+
return &mockConn{}, nil
252+
}
253+
manager := immortalstream.NewManager(logger, dialer)
254+
defer manager.Close()
255+
256+
// Create a test stream
257+
createReq := immortalstream.CreateStreamRequest{TCPPort: 8080}
258+
createResp, err := manager.CreateStream(context.Background(), createReq)
259+
if err != nil {
260+
t.Fatalf("failed to create test stream: %v", err)
261+
}
262+
263+
handler := newImmortalStreamHandler(manager)
264+
265+
// Create chi router to handle URL parameters
266+
r := chi.NewRouter()
267+
r.Get("/{id}", handler.connectToStream)
268+
269+
httpReq := httptest.NewRequest("GET", "/"+createResp.ID.String(), nil)
270+
httpReq.Header.Set("Upgrade", "coder-immortal-stream")
271+
httpReq.Header.Set("Connection", "upgrade")
272+
httpReq.Header.Set("x-coder-immortal-stream-sequence-num", "invalid")
273+
274+
rr := httptest.NewRecorder()
275+
r.ServeHTTP(rr, httpReq)
276+
277+
if rr.Code != http.StatusBadRequest {
278+
t.Fatalf("expected status 400, got %d", rr.Code)
279+
}
280+
})
281+
282+
t.Run("ParseSequenceNumber", func(t *testing.T) {
283+
t.Parallel()
284+
285+
// Test valid sequence number
286+
seqNum, err := parseSequenceNumber("12345")
287+
if err != nil {
288+
t.Fatalf("expected no error, got %v", err)
289+
}
290+
if seqNum != 12345 {
291+
t.Fatalf("expected 12345, got %d", seqNum)
292+
}
293+
294+
// Test empty string
295+
seqNum, err = parseSequenceNumber("")
296+
if err != nil {
297+
t.Fatalf("expected no error for empty string, got %v", err)
298+
}
299+
if seqNum != 0 {
300+
t.Fatalf("expected 0 for empty string, got %d", seqNum)
301+
}
302+
303+
// Test invalid string
304+
_, err = parseSequenceNumber("invalid")
305+
if err == nil {
306+
t.Fatalf("expected error for invalid string")
307+
}
308+
})
309+
310+
t.Run("Routes", func(t *testing.T) {
311+
t.Parallel()
312+
313+
logger := slogtest.Make(t, nil)
314+
manager := immortalstream.NewManager(logger, nil)
315+
defer manager.Close()
316+
317+
handler := newImmortalStreamHandler(manager)
318+
router := handler.Routes()
319+
320+
// Test that routes are properly configured
321+
// This is a basic test to ensure the router doesn't panic
322+
if router == nil {
323+
t.Fatalf("expected non-nil router")
324+
}
325+
})
326+
}
327+
328+
// Mock connection for testing
329+
type mockConn struct {
330+
closed bool
331+
data []byte
332+
pos int
333+
}
334+
335+
func (m *mockConn) Read(b []byte) (int, error) {
336+
if m.closed {
337+
return 0, io.EOF
338+
}
339+
if m.pos >= len(m.data) {
340+
return 0, io.EOF
341+
}
342+
343+
n := copy(b, m.data[m.pos:])
344+
m.pos += n
345+
return n, nil
346+
}
347+
348+
func (m *mockConn) Write(b []byte) (int, error) {
349+
if m.closed {
350+
return 0, xerrors.New("connection closed")
351+
}
352+
m.data = append(m.data, b...)
353+
return len(b), nil
354+
}
355+
356+
func (m *mockConn) Close() error {
357+
m.closed = true
358+
return nil
359+
}
360+
361+
func (m *mockConn) LocalAddr() net.Addr { return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0} }
362+
func (m *mockConn) RemoteAddr() net.Addr { return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 0} }
363+
364+
func (m *mockConn) SetDeadline(t time.Time) error { return nil }
365+
func (m *mockConn) SetReadDeadline(t time.Time) error { return nil }
366+
func (m *mockConn) SetWriteDeadline(t time.Time) error { return nil }

0 commit comments

Comments
 (0)