gpu-turnstile/internal/supervise/supervise_test.go

package supervise

import (
	"context"
	"errors"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
	"time"
)

func TestSplitCommandLine(t *testing.T) {
	cases := []struct {
		in   string
		want []string
	}{
		{`python main.py --port 8188`, []string{"python", "main.py", "--port", "8188"}},
		{`"C:\Program Files\py\python.exe" main.py`, []string{`C:\Program Files\py\python.exe`, "main.py"}},
		{`  spaced   out  `, []string{"spaced", "out"}},
		{`a "b c" d`, []string{"a", "b c", "d"}},
		{"", nil},
	}
	for _, c := range cases {
		got, err := splitCommandLine(c.in)
		if err != nil {
			t.Errorf("splitCommandLine(%q): %v", c.in, err)
			continue
		}
		if len(got) != len(c.want) {
			t.Errorf("splitCommandLine(%q) = %v, want %v", c.in, got, c.want)
			continue
		}
		for i := range got {
			if got[i] != c.want[i] {
				t.Errorf("splitCommandLine(%q) = %v, want %v", c.in, got, c.want)
				break
			}
		}
	}
	if _, err := splitCommandLine(`"unterminated`); err == nil {
		t.Error("expected error for unterminated quote")
	}
}

// TestHelperProcess is the child executed by the lifecycle tests: it just
// sleeps. The env marker is set only around the spawn, so in the normal
// test run this returns immediately.
func TestHelperProcess(t *testing.T) {
	if os.Getenv("GO_HELPER_PROCESS") != "1" {
		return
	}
	time.Sleep(30 * time.Second)
	os.Exit(0)
}

func newHelper(t *testing.T, name string) *Process {
	t.Helper()
	p, err := New(name, `"`+os.Args[0]+`" -test.run=TestHelperProcess`, "",
		func(context.Context) error { return nil }, 5*time.Second, slog.Default())
	if err != nil {
		t.Fatal(err)
	}
	return p
}

// startHelper spawns the child with the marker set; exec.Command inherits
// the environment at spawn time, so it can be unset right after.
func startHelper(t *testing.T, p *Process) {
	t.Helper()
	os.Setenv("GO_HELPER_PROCESS", "1")
	defer os.Unsetenv("GO_HELPER_PROCESS")
	if err := p.EnsureRunning(); err != nil {
		t.Fatal(err)
	}
}

func waitStopped(t *testing.T, p *Process, timeout time.Duration) {
	t.Helper()
	deadline := time.Now().Add(timeout)
	for p.Running() && time.Now().Before(deadline) {
		time.Sleep(50 * time.Millisecond)
	}
	if p.Running() {
		t.Fatal("process still running")
	}
}

func TestEnsureRunningAndStop(t *testing.T) {
	p := newHelper(t, "helper")
	if p.Running() {
		t.Fatal("Running before start")
	}
	startHelper(t, p)
	if !p.Running() {
		t.Fatal("not Running after EnsureRunning")
	}
	if err := p.EnsureRunning(); err != nil {
		t.Fatal("second EnsureRunning must be a no-op")
	}
	p.Stop()
	waitStopped(t, p, 5*time.Second)
}

func TestWaitReady(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	defer srv.Close()
	probe := func(ctx context.Context) error {
		req, _ := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
		resp, err := http.DefaultClient.Do(req)
		if err != nil {
			return err
		}
		resp.Body.Close()
		return nil
	}
	p, err := New("ready", `"`+os.Args[0]+`"`, "", probe, 5*time.Second, slog.Default())
	if err != nil {
		t.Fatal(err)
	}
	if err := p.WaitReady(context.Background()); err != nil {
		t.Fatal(err)
	}

	failing, err := New("failing", `"`+os.Args[0]+`"`, "",
		func(context.Context) error { return errors.New("no") }, 500*time.Millisecond, slog.Default())
	if err != nil {
		t.Fatal(err)
	}
	if err := failing.WaitReady(context.Background()); err == nil {
		t.Fatal("expected timeout error from WaitReady")
	}
}

func TestWatchIdleStops(t *testing.T) {
	p := newHelper(t, "idle")
	startHelper(t, p)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go p.WatchIdle(ctx, 200*time.Millisecond, func() bool { return true })
	waitStopped(t, p, 10*time.Second)
}

func TestWatchIdleRespectsBusyGPU(t *testing.T) {
	p := newHelper(t, "busy")
	startHelper(t, p)
	defer p.Stop()
	ctx, cancel := context.WithCancel(context.Background())
	go p.WatchIdle(ctx, 100*time.Millisecond, func() bool { return false })
	time.Sleep(600 * time.Millisecond)
	cancel()
	if !p.Running() {
		t.Fatal("process was stopped while the GPU was busy")
	}
}