27c8b08a16
Reorganize tests for better structure and coverage: - Move container/security_test.go from internal/ to tests/unit/container/ - Move related tests to proper unit test locations - Delete orphaned test files (startup_blacklist_test.go) - Add privacy middleware unit tests - Add worker config unit tests - Update E2E tests for homelab and websocket scenarios - Update test fixtures with utility functions - Add CLI helper script for arraylist fixes
103 lines
3 KiB
Go
103 lines
3 KiB
Go
package tests
|
|
|
|
import (
|
|
"crypto/tls"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"net/http/httputil"
|
|
"net/url"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
"github.com/jfraeys/fetch_ml/internal/api/datasets"
|
|
"github.com/jfraeys/fetch_ml/internal/api/jobs"
|
|
jupyterj "github.com/jfraeys/fetch_ml/internal/api/jupyter"
|
|
"github.com/jfraeys/fetch_ml/internal/api/ws"
|
|
"github.com/jfraeys/fetch_ml/internal/auth"
|
|
"github.com/jfraeys/fetch_ml/internal/experiment"
|
|
"github.com/jfraeys/fetch_ml/internal/logging"
|
|
)
|
|
|
|
type wsUpgradeProxy struct {
|
|
proxy *httputil.ReverseProxy
|
|
}
|
|
|
|
func (p *wsUpgradeProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
|
// ReverseProxy will forward Upgrade requests; we additionally ensure hop-by-hop
|
|
// headers used for WS are preserved.
|
|
if r.Header.Get("Upgrade") != "" {
|
|
r.Header.Del("Connection")
|
|
r.Header.Add("Connection", "upgrade")
|
|
}
|
|
p.proxy.ServeHTTP(w, r)
|
|
}
|
|
|
|
func startWSBackendServer(t *testing.T) *httptest.Server {
|
|
t.Helper()
|
|
|
|
logger := logging.NewLogger(slog.LevelInfo, false)
|
|
authConfig := &auth.Config{Enabled: false}
|
|
expManager := experiment.NewManager(t.TempDir())
|
|
jobsHandler := jobs.NewHandler(expManager, logger, nil, nil, authConfig, nil)
|
|
jupyterHandler := jupyterj.NewHandler(logger, nil, authConfig)
|
|
datasetsHandler := datasets.NewHandler(logger, nil, "")
|
|
h := ws.NewHandler(authConfig, logger, expManager, "", nil, nil, nil, nil, nil, jobsHandler, jupyterHandler, datasetsHandler)
|
|
|
|
srv := httptest.NewServer(h)
|
|
t.Cleanup(srv.Close)
|
|
return srv
|
|
}
|
|
|
|
func startTLSReverseProxy(t *testing.T, target *url.URL) *httptest.Server {
|
|
t.Helper()
|
|
|
|
rp := httputil.NewSingleHostReverseProxy(target)
|
|
proxyHandler := &wsUpgradeProxy{proxy: rp}
|
|
|
|
proxySrv := httptest.NewTLSServer(proxyHandler)
|
|
return proxySrv
|
|
}
|
|
|
|
func TestWSS_UpgradeThroughTLSReverseProxy(t *testing.T) {
|
|
backendSrv := startWSBackendServer(t)
|
|
backendURL, err := url.Parse(backendSrv.URL)
|
|
if err != nil {
|
|
t.Fatalf("failed to parse backend url: %v", err)
|
|
}
|
|
|
|
proxySrv := startTLSReverseProxy(t, backendURL)
|
|
defer proxySrv.Close()
|
|
|
|
proxyURL, err := url.Parse(proxySrv.URL)
|
|
if err != nil {
|
|
t.Fatalf("failed to parse proxy url: %v", err)
|
|
}
|
|
|
|
wssURL := url.URL{Scheme: "wss", Host: proxyURL.Host, Path: "/ws"}
|
|
|
|
dialer := websocket.Dialer{
|
|
TLSClientConfig: &tls.Config{
|
|
MinVersion: tls.VersionTLS12,
|
|
InsecureSkipVerify: true, // test-only (self-signed cert from httptest)
|
|
},
|
|
}
|
|
|
|
conn, resp, err := dialer.Dial(wssURL.String(), nil)
|
|
if resp != nil && resp.Body != nil {
|
|
defer func() { _ = resp.Body.Close() }()
|
|
}
|
|
if err != nil {
|
|
t.Fatalf("failed to connect via wss through proxy: %v", err)
|
|
}
|
|
defer func() { _ = conn.Close() }()
|
|
|
|
// Basic write to ensure upgraded channel is usable.
|
|
_ = conn.SetWriteDeadline(time.Now().Add(5 * time.Second))
|
|
statusMsg := []byte{0x02}
|
|
statusMsg = append(statusMsg, make([]byte, 16)...)
|
|
if err := conn.WriteMessage(websocket.BinaryMessage, statusMsg); err != nil {
|
|
t.Fatalf("failed to write websocket message: %v", err)
|
|
}
|
|
}
|