277 lines
8.0 KiB
Go
277 lines
8.0 KiB
Go
package supervise
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"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()
|
|
// The probe always fails: nothing external serves the port, so
|
|
// EnsureRunning spawns the helper child.
|
|
p, err := New(name, `"`+os.Args[0]+`" -test.run=TestHelperProcess`, "",
|
|
func(context.Context) error { return errors.New("nothing there") }, 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 TestEnsureRunningPrefersExternalServer(t *testing.T) {
|
|
// The port is already served (e.g. the ComfyUI desktop app): no child
|
|
// is spawned, the supervisor reports ready, and Stop is a no-op.
|
|
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("external", `"`+os.Args[0]+`"`, "", probe, 5*time.Second, slog.Default())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := p.EnsureRunning(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if p.Running() {
|
|
t.Fatal("spawned a child even though the port is already served")
|
|
}
|
|
if !p.Ready() {
|
|
t.Fatal("external server should count as ready")
|
|
}
|
|
p.Stop() // must not touch the external server
|
|
}
|
|
|
|
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")
|
|
}
|
|
}
|
|
|
|
func TestComfyLayoutAndDefaultCommand(t *testing.T) {
|
|
// Nested layout (ComfyUI/main.py under dir) wins.
|
|
dir := t.TempDir()
|
|
nested := filepath.Join(dir, "ComfyUI", "main.py")
|
|
if err := os.MkdirAll(filepath.Dir(nested), 0o755); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := os.WriteFile(nested, []byte("x"), 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
python, script := ComfyLayout("windows", dir)
|
|
if want := filepath.Join(dir, ".venv", "Scripts", "python.exe"); python != want {
|
|
t.Errorf("python = %s, want %s", python, want)
|
|
}
|
|
if script != nested {
|
|
t.Errorf("script = %s, want %s", script, nested)
|
|
}
|
|
cmd := DefaultComfyCommand("windows", dir, "http://127.0.0.1:8189")
|
|
want := `"` + filepath.Join(dir, ".venv", "Scripts", "python.exe") + `" ` + filepath.Join("ComfyUI", "main.py") + " --port 8189"
|
|
if cmd != want {
|
|
t.Errorf("cmd = %q, want %q", cmd, want)
|
|
}
|
|
|
|
// Flat layout (main.py directly under dir) is found too.
|
|
flat := t.TempDir()
|
|
if err := os.WriteFile(filepath.Join(flat, "main.py"), []byte("x"), 0o644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, script := ComfyLayout("linux", flat); script != filepath.Join(flat, "main.py") {
|
|
t.Errorf("flat script = %s", script)
|
|
}
|
|
cmd = DefaultComfyCommand("linux", flat, "http://comfy.internal")
|
|
if strings.Contains(cmd, "--port") {
|
|
t.Errorf("cmd = %q, want no --port for a port-less URL", cmd)
|
|
}
|
|
if !strings.HasSuffix(cmd, `" main.py`) {
|
|
t.Errorf("cmd = %q, want quoted python + relative main.py", cmd)
|
|
}
|
|
|
|
// Neither exists yet: default to the documented nested form so the
|
|
// startup warning points there.
|
|
empty := t.TempDir()
|
|
if _, script := ComfyLayout("windows", empty); script != filepath.Join(empty, "ComfyUI", "main.py") {
|
|
t.Errorf("missing-layout script = %s", script)
|
|
}
|
|
}
|
|
|
|
func TestStripANSI(t *testing.T) {
|
|
cases := map[string]string{
|
|
"\x1b[32m[INFO]\x1b[0m Starting server": "[INFO] Starting server",
|
|
"\x1b[1m\x1b[31m[ERROR]\x1b[0m boom": "[ERROR] boom",
|
|
"plain line": "plain line",
|
|
"\x1b[33mWARN\x1b[0m: \x1b[1mbold\x1b[0m": "WARN: bold",
|
|
}
|
|
for in, want := range cases {
|
|
if got := stripANSI(in); got != want {
|
|
t.Errorf("stripANSI(%q) = %q, want %q", in, got, want)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestDesktopSharedDir(t *testing.T) {
|
|
root := t.TempDir()
|
|
comfy := filepath.Join(root, "ComfyUI-Installs", "rtx5080", "ComfyUI")
|
|
if err := os.MkdirAll(comfy, 0o755); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if got := DesktopSharedDir(comfy); got != "" {
|
|
t.Fatalf("no shared dir yet: got %q, want empty", got)
|
|
}
|
|
shared := filepath.Join(root, "ComfyUI-Shared")
|
|
if err := os.MkdirAll(filepath.Join(shared, "models"), 0o755); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if got := DesktopSharedDir(comfy); got != shared {
|
|
t.Fatalf("got %q, want %q", got, shared)
|
|
}
|
|
if got := DesktopSharedDir(filepath.Join(root, "plain", "ComfyUI")); got != "" {
|
|
t.Fatalf("non-desktop layout: got %q, want empty", got)
|
|
}
|
|
}
|