4 files changed, 28 insertions, 3 deletions
diff --git a/internal/handlers/handlers.go b/internal/handlers/handlers.go
index c44e771..b8fc574 100644
--- a/internal/handlers/handlers.go
+++ b/internal/handlers/handlers.go
@@ -207,6 +207,11 @@ func (h *Handler) aggregateData(ctx context.Context, forceRefresh bool) (*models
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
 		boards, err := h.fetchBoards(ctx, forceRefresh)
 		mu.Lock()
 		defer mu.Unlock()
@@ -221,6 +226,11 @@ func (h *Handler) aggregateData(ctx context.Context, forceRefresh bool) (*models
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
 		tasks, err := h.fetchTasks(ctx, forceRefresh)
 		mu.Lock()
 		defer mu.Unlock()
@@ -255,6 +265,11 @@ func (h *Handler) aggregateData(ctx context.Context, forceRefresh bool) (*models
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
 		projects, err := h.todoistClient.GetProjects(ctx)
 		mu.Lock()
 		defer mu.Unlock()
@@ -270,6 +285,11 @@ func (h *Handler) aggregateData(ctx context.Context, forceRefresh bool) (*models
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
+			select {
+			case <-ctx.Done():
+				return
+			default:
+			}
 			meals, err := h.fetchMeals(ctx, forceRefresh)
 			mu.Lock()
 			defer mu.Unlock()
@@ -286,6 +306,11 @@ func (h *Handler) aggregateData(ctx context.Context, forceRefresh bool) (*models
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
+			select {
+			case <-ctx.Done():
+				return
+			default:
+			}
 			events, err := h.googleCalendarClient.GetUpcomingEvents(ctx, 10)
 			mu.Lock()
 			defer mu.Unlock()
diff --git a/internal/handlers/handlers_test.go b/internal/handlers/handlers_test.go
index e4a9f05..5628237 100644
--- a/internal/handlers/handlers_test.go
+++ b/internal/handlers/handlers_test.go
@@ -39,7 +39,7 @@ func setupTestDB(t *testing.T) (*store.Store, func()) {
 	}
 
 	// Initialize store (this runs migrations)
-	db, err := store.New(tmpFile.Name())
+	db, err := store.New(tmpFile.Name(), "migrations")
 	if err != nil {
 		os.Chdir(originalDir)
 		os.Remove(tmpFile.Name())
diff --git a/internal/handlers/heuristic_test.go b/internal/handlers/heuristic_test.go
index 2b70218..b03b664 100644
--- a/internal/handlers/heuristic_test.go
+++ b/internal/handlers/heuristic_test.go
@@ -33,7 +33,7 @@ func TestHandleTasks_Heuristic(t *testing.T) {
 	defer os.Chdir(originalDir)
 
 	// Initialize store (this runs migrations)
-	db, err := store.New(tmpFile.Name())
+	db, err := store.New(tmpFile.Name(), "migrations")
 	if err != nil {
 		t.Fatalf("Failed to initialize store: %v", err)
 	}
diff --git a/internal/handlers/tab_state_test.go b/internal/handlers/tab_state_test.go
index d7bb8dd..60f1340 100644
--- a/internal/handlers/tab_state_test.go
+++ b/internal/handlers/tab_state_test.go
@@ -13,7 +13,7 @@ import (
 
 func TestHandleDashboard_TabState(t *testing.T) {
 	// Create a temporary database for testing
-	db, err := store.New(":memory:")
+	db, err := store.New(":memory:", "../../migrations")
 	if err != nil {
 		t.Fatalf("Failed to create test database: %v", err)
 	}