dbf96020af
Phase 1: Fix Redis Schema Leak - Create internal/storage/dataset.go with DatasetStore abstraction - Remove all direct Redis calls from cmd/data_manager/data_sync.go - data_manager now uses DatasetStore for transfer tracking and metadata Phase 2: Simplify TUI Services - Embed *queue.TaskQueue directly in services.TaskQueue - Eliminate 60% of wrapper boilerplate (203 -> ~100 lines) - Keep only TUI-specific methods (EnqueueTask, GetJobStatus, experiment methods) Phase 5: Clean go.mod Dependencies - Remove duplicate go-redis/redis/v8 dependency - Migrate internal/storage/migrate.go to redis/go-redis/v9 - Separate test-only deps (miniredis, testify) into own block Results: - Zero direct Redis calls in cmd/ - 60% fewer lines in TUI services - Cleaner dependency structure
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)
|
|
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)
|
|
}
|
|
}
|