cost-model/pkg/cmd/costmodel/costmodel_test.go

package costmodel

import (
	"context"
	"fmt"
	"net/http"
	"testing"
	"time"

	"github.com/opencost/opencost/pkg/costmodel"
	"github.com/opencost/opencost/pkg/env"
)

func TestMCPServerGracefulShutdown(t *testing.T) {
	// Test that MCP server responds to context cancellation and shuts down gracefully

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	accesses := &costmodel.Accesses{}
	port := env.GetMCPHTTPPort()

	// Channel to signal when server is ready
	serverReady := make(chan error, 1)

	// Start MCP server
	go func() {
		err := StartMCPServer(ctx, accesses, nil)
		serverReady <- err
	}()

	// Wait for server to be ready by attempting to connect
	serverUp := false
	for i := 0; i < 10; i++ {
		time.Sleep(100 * time.Millisecond)
		client := &http.Client{Timeout: 1 * time.Second}
		resp, err := client.Get(fmt.Sprintf("http://localhost:%d/", port))
		if err == nil {
			resp.Body.Close()
			serverUp = true
			break
		}
	}

	if !serverUp {
		t.Skip("MCP server did not start (may be expected in test environment)")
	}

	// Trigger shutdown by cancelling context
	shutdownStart := time.Now()
	cancel()

	// Wait for shutdown to complete (with reasonable timeout)
	shutdownDone := make(chan bool, 1)
	go func() {
		time.Sleep(15 * time.Second)
		shutdownDone <- false
	}()

	// Give shutdown goroutine time to execute
	time.Sleep(1 * time.Second)

	// Verify server is no longer accepting connections
	client := &http.Client{Timeout: 1 * time.Second}
	_, err := client.Get(fmt.Sprintf("http://localhost:%d/", port))
	if err == nil {
		t.Error("Server still accepting connections after shutdown")
	}

	shutdownDone <- true
	<-shutdownDone

	shutdownDuration := time.Since(shutdownStart)
	t.Logf("Graceful shutdown completed in %v", shutdownDuration)

	// Verify shutdown completed in reasonable time (should be much less than 12s)
	if shutdownDuration > 12*time.Second {
		t.Errorf("Shutdown took too long: %v (expected < 12s)", shutdownDuration)
	}
}