From 2a922542b183f61339a31d80517c7ff1314ee462 Mon Sep 17 00:00:00 2001
From: Jeremie Fraeys <jfaeys@gmail.com>
Date: Tue, 17 Feb 2026 20:51:57 -0500
Subject: [PATCH] test: fix ws_test.go to use updated NewHandler signature

Updated test file to pass jobs, jupyter, and datasets handlers to NewHandler.
All tests now pass.
---
 tests/unit/api/ws_test.go | 39 +++++++++++++++++++++++++++------------
 1 file changed, 27 insertions(+), 12 deletions(-)

diff --git a/tests/unit/api/ws_test.go b/tests/unit/api/ws_test.go
index bb494aa..cc24c11 100644
--- a/tests/unit/api/ws_test.go
+++ b/tests/unit/api/ws_test.go
@@ -8,6 +8,9 @@ import (
 	"strings"
 	"testing"
 
+	"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"
 	wspkg "github.com/jfraeys/fetch_ml/internal/api/ws"
 	"github.com/jfraeys/fetch_ml/internal/auth"
 	"github.com/jfraeys/fetch_ml/internal/experiment"
@@ -15,13 +18,16 @@ import (
 )
 
 func TestNewWSHandler(t *testing.T) {
-	t.Parallel() // Enable parallel execution
+	t.Parallel()
 
 	authConfig := &auth.Config{}
-	logger := logging.NewLogger(slog.LevelInfo, false) // Create a real logger
+	logger := logging.NewLogger(slog.LevelInfo, false)
 	expManager := experiment.NewManager("/tmp")
+	jobsHandler := jobs.NewHandler(expManager, logger, nil, nil, authConfig)
+	jupyterHandler := jupyterj.NewHandler(logger, nil, authConfig)
+	datasetsHandler := datasets.NewHandler(logger, nil, "")
 
-	handler := wspkg.NewHandler(authConfig, logger, expManager, "", nil, nil, nil, nil, nil)
+	handler := wspkg.NewHandler(authConfig, logger, expManager, "", nil, nil, nil, nil, nil, jobsHandler, jupyterHandler, datasetsHandler)
 
 	if handler == nil {
 		t.Error("Expected non-nil WSHandler")
@@ -50,13 +56,16 @@ func TestWSHandlerConstants(t *testing.T) {
 }
 
 func TestWSHandlerWebSocketUpgrade(t *testing.T) {
-	t.Parallel() // Enable parallel execution
+	t.Parallel()
 
 	authConfig := &auth.Config{}
-	logger := logging.NewLogger(slog.LevelInfo, false) // Create a real logger
+	logger := logging.NewLogger(slog.LevelInfo, false)
 	expManager := experiment.NewManager("/tmp")
+	jobsHandler := jobs.NewHandler(expManager, logger, nil, nil, authConfig)
+	jupyterHandler := jupyterj.NewHandler(logger, nil, authConfig)
+	datasetsHandler := datasets.NewHandler(logger, nil, "")
 
-	handler := wspkg.NewHandler(authConfig, logger, expManager, "", nil, nil, nil, nil, nil)
+	handler := wspkg.NewHandler(authConfig, logger, expManager, "", nil, nil, nil, nil, nil, jobsHandler, jupyterHandler, datasetsHandler)
 
 	// Create a test HTTP request
 	req := httptest.NewRequest("GET", "/ws", nil)
@@ -87,13 +96,16 @@ func TestWSHandlerWebSocketUpgrade(t *testing.T) {
 }
 
 func TestWSHandlerInvalidRequest(t *testing.T) {
-	t.Parallel() // Enable parallel execution
+	t.Parallel()
 
 	authConfig := &auth.Config{}
-	logger := logging.NewLogger(slog.LevelInfo, false) // Create a real logger
+	logger := logging.NewLogger(slog.LevelInfo, false)
 	expManager := experiment.NewManager("/tmp")
+	jobsHandler := jobs.NewHandler(expManager, logger, nil, nil, authConfig)
+	jupyterHandler := jupyterj.NewHandler(logger, nil, authConfig)
+	datasetsHandler := datasets.NewHandler(logger, nil, "")
 
-	handler := wspkg.NewHandler(authConfig, logger, expManager, "", nil, nil, nil, nil, nil)
+	handler := wspkg.NewHandler(authConfig, logger, expManager, "", nil, nil, nil, nil, nil, jobsHandler, jupyterHandler, datasetsHandler)
 
 	// Create a test HTTP request without WebSocket headers
 	req := httptest.NewRequest("GET", "/ws", nil)
@@ -112,13 +124,16 @@ func TestWSHandlerInvalidRequest(t *testing.T) {
 }
 
 func TestWSHandlerPostRequest(t *testing.T) {
-	t.Parallel() // Enable parallel execution
+	t.Parallel()
 
 	authConfig := &auth.Config{}
-	logger := logging.NewLogger(slog.LevelInfo, false) // Create a real logger
+	logger := logging.NewLogger(slog.LevelInfo, false)
 	expManager := experiment.NewManager("/tmp")
+	jobsHandler := jobs.NewHandler(expManager, logger, nil, nil, authConfig)
+	jupyterHandler := jupyterj.NewHandler(logger, nil, authConfig)
+	datasetsHandler := datasets.NewHandler(logger, nil, "")
 
-	handler := wspkg.NewHandler(authConfig, logger, expManager, "", nil, nil, nil, nil, nil)
+	handler := wspkg.NewHandler(authConfig, logger, expManager, "", nil, nil, nil, nil, nil, jobsHandler, jupyterHandler, datasetsHandler)
 
 	// Create a POST request (should fail)
 	req := httptest.NewRequest("POST", "/ws", strings.NewReader("data"))