go-omni/cmd/go-omni/apigen/testdata/managed/phases_test.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"example.com/managed/internal/svc"
	"git.yuntuo.hk/k2safe/go-omni/lifecycle"
	"io"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func phasedConfig(t *testing.T, address, mode string, fail bool) string {
	t.Helper()
	path := filepath.Join(t.TempDir(), "config.json")
	drainTimeout := "1s"
	if mode == "stuck" {
		drainTimeout = "50ms"
	}
	data, _ := json.Marshal(map[string]any{"backgroundMode": mode, "failInit": fail, "server": map[string]any{"address": address, "transport": "${GO_OMNI_TRANSPORT:-nethttp}", "lifecycle": map[string]string{"stopTimeout": "1s", "drainTimeout": drainTimeout, "cancelTimeout": "50ms"}}})
	if err := os.WriteFile(path, data, 0600); err != nil {
		t.Fatal(err)
	}
	return path
}
func await(t *testing.T, ch <-chan struct{}) {
	t.Helper()
	select {
	case <-ch:
	case <-time.After(3 * time.Second):
		t.Fatal("phase did not arrive")
	}
}

func TestGeneratedBackgroundPhases(t *testing.T) {
	for _, transport := range []string{"nethttp", "native"} {
		t.Run(transport, func(t *testing.T) {
			t.Setenv("GO_OMNI_TRANSPORT", transport)
			listener, err := net.Listen("tcp", "127.0.0.1:0")
			if err != nil {
				t.Fatal(err)
			}
			addr := listener.Addr().String()
			_ = listener.Close()
			svc.Started = make(chan struct{}, 1)
			svc.Release = make(chan struct{})
			before := svc.Closed.Load()
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()
			done := make(chan error, 1)
			path := phasedConfig(t, addr, "hold", false)
			go func() { done <- run(ctx, path) }()
			waitReady(t, "http://"+addr, done)
			state := svc.Background.Load()
			defer func() {
				select {
				case <-state.Release:
				default:
					close(state.Release)
				}
			}()
			requestDone := make(chan error, 1)
			go func() {
				client := http.Client{Timeout: 3 * time.Second}
				resp, e := client.Get("http://" + addr + "/slow")
				if e == nil {
					_, _ = io.Copy(io.Discard, resp.Body)
					_ = resp.Body.Close()
					if resp.StatusCode != 200 {
						e = errors.New("HTTP drain failed")
					}
				}
				requestDone <- e
			}()
			await(t, svc.Started)
			deps := svc.Current.Load()
			deps.RequestShutdown(svc.FatalDependency)
			if deps.WorkContext().Err() != nil {
				t.Fatal("accepted work canceled at signal")
			}
			select {
			case <-state.Stop:
				t.Fatal("stopped background before HTTP finished")
			default:
			}
			close(svc.Release)
			if err := <-requestDone; err != nil {
				t.Fatal(err)
			}
			await(t, state.Drain)
			if svc.Closed.Load() != before {
				t.Fatal("client closed during background drain")
			}
			close(state.Release)
			if err := <-done; !errors.Is(err, svc.FatalDependency) || errors.Is(err, lifecycle.ErrResourcesRetained) {
				t.Fatal(err)
			}
			if !state.Wrote.Load() || svc.Closed.Load() != before+1 {
				t.Fatal("missing final write or cleanup")
			}
			if err := state.Future.Wait(context.Background()); err != nil {
				t.Fatal(err)
			}
		})
	}
}

func TestGeneratedBackgroundFailureCleanup(t *testing.T) {
	for _, failInit := range []bool{false, true} {
		t.Run(map[bool]string{true: "initialization", false: "listen"}[failInit], func(t *testing.T) {
			l, err := net.Listen("tcp", "127.0.0.1:0")
			if err != nil {
				t.Fatal(err)
			}
			defer l.Close()
			before := svc.Closed.Load()
			err = run(context.Background(), phasedConfig(t, l.Addr().String(), "auto", failInit))
			if err == nil || errors.Is(err, lifecycle.ErrResourcesRetained) {
				t.Fatal(err)
			}
			state := svc.Background.Load()
			if !state.Wrote.Load() || svc.Closed.Load() != before+1 {
				t.Fatal("failure cleanup did not join background first")
			}
		})
	}
}

func TestGeneratedBackgroundTimeoutProcess(t *testing.T) {
	if path := os.Getenv("OMNI_PHASE_TIMEOUT_CONFIG"); path != "" {
		os.Args = []string{os.Args[0], "-config", path}
		main()
		return
	}
	path := phasedConfig(t, "127.0.0.1:0", "stuck", true)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	cmd := exec.CommandContext(ctx, os.Args[0], "-test.run=^TestGeneratedBackgroundTimeoutProcess$")
	cmd.Env = append(os.Environ(), "OMNI_PHASE_TIMEOUT_CONFIG="+path)
	output, err := cmd.CombinedOutput()
	var exit *exec.ExitError
	if !errors.As(err, &exit) || exit.ExitCode() != 1 || ctx.Err() != nil {
		t.Fatal(err, string(output))
	}
	if strings.Contains(string(output), "test.client.closed") || !strings.Contains(string(output), "shared resources retained") || !strings.Contains(string(output), "client initialization failed") {
		t.Fatal(string(output))
	}
}