package api import ( "net/http" "net/http/httptest" "strings" "testing" "time" "github.com/gorilla/websocket" ) func TestJobConnection_ConnectAndClose(t *testing.T) { upgrader := websocket.Upgrader{} server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { conn, err := upgrader.Upgrade(w, r, nil) if err != nil { return } defer conn.Close() var msg map[string]interface{} if err := conn.ReadJSON(&msg); err != nil { return } if msg["type"] == "auth" { _ = conn.WriteJSON(map[string]string{"type": "auth_ok"}) } // Keep open briefly so client can mark connected. time.Sleep(100 * time.Millisecond) })) defer server.Close() jc := NewJobConnection() managerURL := strings.Replace(server.URL, "http://", "http://", 1) if err := jc.Connect(managerURL, "/job/1", "token123"); err != nil { t.Fatalf("Connect failed: %v", err) } if !jc.IsConnected() { t.Fatal("expected connection to be marked connected") } jc.Close() }