94 lines
2.9 KiB
Go
94 lines
2.9 KiB
Go
package observability
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/coder/websocket"
|
|
)
|
|
|
|
func TestWebSocketHandshakeIsNotConnectionLifetime(t *testing.T) {
|
|
r := NewRegistry()
|
|
released := make(chan struct{})
|
|
done := make(chan struct{})
|
|
h := r.Middleware(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
|
|
defer close(done)
|
|
c, err := websocket.Accept(w, req, nil)
|
|
if err != nil {
|
|
return
|
|
}
|
|
defer c.CloseNow()
|
|
<-released
|
|
}))
|
|
server := httptest.NewTLSServer(h)
|
|
defer server.Close()
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
c, _, err := websocket.Dial(ctx, strings.Replace(server.URL, "https://", "wss://", 1)+"/v1/worker/connect", &websocket.DialOptions{HTTPClient: server.Client()})
|
|
if err != nil {
|
|
close(released)
|
|
t.Fatal(err)
|
|
}
|
|
defer c.CloseNow()
|
|
// Dial returns after the upgrade response, while the server handler stays open.
|
|
deadline := time.Now().Add(time.Second)
|
|
for !strings.Contains(r.Render(), `orbit_websocket_connections{route="worker_connect"} 1`) {
|
|
if time.Now().After(deadline) {
|
|
close(released)
|
|
t.Fatal("upgrade not observed")
|
|
}
|
|
time.Sleep(time.Millisecond)
|
|
}
|
|
out := r.Render()
|
|
close(released)
|
|
<-done
|
|
if !strings.Contains(out, `orbit_http_requests_total{method="GET",route="worker_connect",status_class="1xx"} 1`) || !strings.Contains(out, "orbit_http_requests_in_flight 0\n") {
|
|
t.Fatal("upgrade counted as unfinished HTTP")
|
|
}
|
|
// done closes inside the handler; wait for middleware's defer to publish closure.
|
|
deadline = time.Now().Add(time.Second)
|
|
for !strings.Contains(r.Render(), `orbit_websocket_connections{route="worker_connect"} 0`) {
|
|
if time.Now().After(deadline) {
|
|
t.Fatal("connection gauge leaked")
|
|
}
|
|
time.Sleep(time.Millisecond)
|
|
}
|
|
if !strings.Contains(r.Render(), `orbit_http_requests_total{method="GET",route="worker_connect",status_class="1xx"} 1`) {
|
|
t.Fatal("upgrade counted twice")
|
|
}
|
|
}
|
|
func TestInformationalResponseAndStreamingPreserveHTTP(t *testing.T) {
|
|
r := NewRegistry()
|
|
server := httptest.NewServer(r.Middleware(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
|
|
w.WriteHeader(103)
|
|
w.WriteHeader(201)
|
|
_, _ = io.WriteString(w, "first")
|
|
if err := http.NewResponseController(w).Flush(); err != nil {
|
|
panic(err)
|
|
}
|
|
_, _ = io.WriteString(w, "second")
|
|
})))
|
|
defer server.Close()
|
|
resp, err := server.Client().Get(server.URL + "/v1/tasks")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer resp.Body.Close()
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil || resp.StatusCode != 201 || string(body) != "firstsecond" {
|
|
t.Fatal("wrapper changed HTTP response", resp.StatusCode, err)
|
|
}
|
|
deadline := time.Now().Add(time.Second)
|
|
for !strings.Contains(r.Render(), `route="tasks",status_class="2xx"} 1`) {
|
|
if time.Now().After(deadline) {
|
|
t.Fatal("103 blocked final status metric")
|
|
}
|
|
time.Sleep(time.Millisecond)
|
|
}
|
|
}
|