myxce/src/frontend/xcecli/cmd/jobs/submit_test.go

package jobs

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	"go.uber.org/zap"

	"xcecli/common/format"
	"xcecli/common/request"
	"xcecli/config"
	"xcecli/models"
)

// bug-1676: BuildJobRequest asserts our own hostname as FromHost rather than
// leaving it for the server to derive from the connection (spoofable via
// X-Forwarded-For, and only ever an IP, never a hostname).
func TestBuildJobRequest_SetsFromHostToLocalHostname(t *testing.T) {
	js := NewJobSubmitter(nil)
	js.command = "sleep 1"

	req, err := js.BuildJobRequest()
	if err != nil {
		t.Fatalf("BuildJobRequest: %v", err)
	}

	wantHost, hostErr := os.Hostname()
	if hostErr != nil {
		t.Skipf("os.Hostname() unavailable in this environment: %v", hostErr)
	}
	if req.FromHost != wantHost {
		t.Errorf("FromHost = %q, want local hostname %q", req.FromHost, wantHost)
	}
}

// submitServer answers every submit with the given response data.
func submitServer(t *testing.T, data models.JobSubmitData) *httptest.Server {
	t.Helper()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(models.JobSubmitResponse{Code: 0, Msg: "ok", Data: data})
	}))
	t.Cleanup(srv.Close)
	// The submit path logs through config.Log, which only main initialises.
	if config.Log == nil {
		config.Log = zap.NewNop()
		t.Cleanup(func() { config.Log = nil })
	}
	return srv
}

// An accepted array submission prints the array identity in written form,
// without an element number, and the server's element count (3.1, 3.2); the
// count is mirrored in the client model for scripted use.
func TestSubmitOutput_ArrayIdentityAndCount(t *testing.T) {
	srv := submitServer(t, models.JobSubmitData{
		JobID:     int64(testArrayKey),
		Message:   fmt.Sprintf("Job <%d> is submitted to queue <normal> (3 elements).", testArrayKey),
		ArraySize: 3,
	})
	js := NewJobSubmitter(nil)
	js.command = "sleep 1"
	req, err := js.BuildJobRequest()
	if err != nil {
		t.Fatalf("BuildJobRequest: %v", err)
	}
	var resp *models.JobSubmitResponse
	out, err := captureStdout(t, func() error {
		var e error
		resp, e = submitAndReport(request.NewAPIClient(srv.URL), "tok", req)
		return e
	})
	if err != nil {
		t.Fatalf("submitAndReport: %v", err)
	}
	want := "Job <" + testArrayWritten + "> is submitted to queue <normal> (3 elements).\n"
	if out != want {
		t.Fatalf("output = %q, want %q", out, want)
	}
	if resp.Data.ArraySize != 3 {
		t.Errorf("ArraySize = %d, want 3", resp.Data.ArraySize)
	}
}

// An ordinary submission prints exactly as before (3.3): no count added.
func TestSubmitOutput_OrdinaryUnchanged(t *testing.T) {
	id := format.MakeJobID(0, 0, 1700000000, 0, 7)
	srv := submitServer(t, models.JobSubmitData{
		JobID:   int64(id),
		Message: fmt.Sprintf("Job <%d> is submitted to queue <normal>.", id),
	})
	js := NewJobSubmitter(nil)
	js.command = "sleep 1"
	req, err := js.BuildJobRequest()
	if err != nil {
		t.Fatalf("BuildJobRequest: %v", err)
	}
	var resp *models.JobSubmitResponse
	out, err := captureStdout(t, func() error {
		var e error
		resp, e = submitAndReport(request.NewAPIClient(srv.URL), "tok", req)
		return e
	})
	if err != nil {
		t.Fatalf("submitAndReport: %v", err)
	}
	want := "Job <0-0-1700000000-7> is submitted to queue <normal>.\n"
	if out != want {
		t.Fatalf("output = %q, want %q", out, want)
	}
	if resp.Data.ArraySize != 0 {
		t.Errorf("ArraySize = %d, want 0", resp.Data.ArraySize)
	}
}

// Submit help describes the array forms, numbering from 1 and the limit (9.1).
func TestSubmitHelpShowsArrayForms(t *testing.T) {
	long := NewSubmitJobsCmd(nil).Long
	for _, s := range []string{
		"-J 'name[1-N]'", "-J 'name[1]'", "start at 1",
		"set by the API server", "default 1000", "32767",
	} {
		if !strings.Contains(long, s) {
			t.Errorf("submit help lacks %q", s)
		}
	}
}

// getEnvironmentVariables must keep entries whose value contains '=' (the
// apiserver splits at the FIRST '=', so A=b=c is A -> "b=c") -- dropping
// them used to be silent data loss. A nameless or '='-less entry cannot
// become a map key, so it is warned about and dropped.
func TestGetEnvironmentVariables_KeepsMultiEqualSign(t *testing.T) {
	js := NewJobSubmitter(nil)
	js.subEnvVars = []string{"A=1", "B=b=c", "C=", "=novalue", "novalue"}

	r, w, pipeErr := os.Pipe()
	if pipeErr != nil {
		t.Fatalf("pipe: %v", pipeErr)
	}
	old := os.Stdout
	os.Stdout = w
	done := make(chan string)
	go func() {
		var buf bytes.Buffer
		_, _ = io.Copy(&buf, r)
		done <- buf.String()
	}()
	got := js.getEnvironmentVariables()
	os.Stdout = old
	_ = w.Close()
	warn := <-done

	want := []string{"A=1", "B=b=c", "C="}
	if len(got) != len(want) {
		t.Fatalf("env = %q, want %q", got, want)
	}
	for i := range want {
		if got[i] != want[i] {
			t.Errorf("env[%d] = %q, want %q", i, got[i], want[i])
		}
	}
	if !strings.Contains(warn, "=novalue") || !strings.Contains(warn, "warning:") {
		t.Errorf("warning output = %q, want it to name the nameless entry", warn)
	}
	if !strings.Contains(warn, "novalue") {
		t.Errorf("warning output = %q, want it to name the '='-less entry", warn)
	}
}

// Attach marks the terminal-attach kind of TTY job so the worker may reap
// ttyd when no client ever connects; --browse keeps the web-viewed
// semantics and must not set it.
func TestBuildJobRequest_AttachFlagMatrix(t *testing.T) {
	cases := []struct {
		name        string
		interactive bool
		useTTY      bool
		useBrowse   bool
		useXJob     bool
		wantInter   int
		wantAttach  int
	}{
		{"batch job", false, false, false, false, 0, 0},
		{"tty attach (bare -I)", true, true, false, false, 1, 1},
		{"tty browse", true, true, true, false, 1, 0},
		{"xjob", true, false, false, true, 2, 0},
		{"display (interactive=3)", true, false, false, false, 3, 0},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			js := NewJobSubmitter(nil)
			js.command = "sleep 1"
			js.interactive = tc.interactive
			js.useTTY = tc.useTTY
			js.useBrowse = tc.useBrowse
			js.useXJob = tc.useXJob
			req, err := js.BuildJobRequest()
			if err != nil {
				t.Fatalf("BuildJobRequest: %v", err)
			}
			if req.Interactive != tc.wantInter {
				t.Errorf("Interactive = %d, want %d", req.Interactive, tc.wantInter)
			}
			if req.Attach != tc.wantAttach {
				t.Errorf("Attach = %d, want %d", req.Attach, tc.wantAttach)
			}
		})
	}
}

// jobStateServer answers the per-job status query with the states fed in
// sequence, holding the last one forever; every request's arrival time is
// recorded so tests can assert the poll cadence.
type jobStateServer struct {
	srv     *httptest.Server
	states  chan string
	mu      sync.Mutex
	pollAts []time.Time
	// urlQueries counts requests carrying fields=interactive_job_url.
	urlQueries int
}

func newJobStateServer(t *testing.T, states ...string) *jobStateServer {
	t.Helper()
	js := &jobStateServer{states: make(chan string, len(states))}
	for _, s := range states {
		js.states <- s
	}
	held := ""
	js.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		js.mu.Lock()
		js.pollAts = append(js.pollAts, time.Now())
		if r.URL.Query().Get("fields") == "interactive_job_url" {
			js.urlQueries++
		}
		js.mu.Unlock()
		select {
		case held = <-js.states:
		default:
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(models.JobsResponse{
			Code: 0,
			Data: struct {
				Count      int               `json:"count"`
				Rows       []models.QueryJob `json:"rows"`
				Page_token string            `json:"page_token"`
			}{Rows: []models.QueryJob{{JobStatus: held, InteractiveJobURL: "http://h:9000"}}},
		})
	}))
	t.Cleanup(js.srv.Close)
	return js
}

func (js *jobStateServer) pollCount() int {
	js.mu.Lock()
	defer js.mu.Unlock()
	return len(js.pollAts)
}

// waitForInteractiveURL fast-polls the PEND phase (500ms) so a normally
// dispatched job attaches within ~a second, then backs off to 2s so a job
// waiting on a long queue doesn't hold 2 qps against the server forever.
func TestWaitForInteractiveURL_PendFastThenBackoff(t *testing.T) {
	if config.Log == nil {
		config.Log = zap.NewNop()
		t.Cleanup(func() { config.Log = nil })
	}
	// PEND through the fast window (11 polls at t=0..5.0s) plus one backoff
	// round, then RUN+URL: the total wait is 10*0.5 + 2 + ~0 = ~7s.
	states := make([]string, 11)
	for i := range states {
		states[i] = "PEND"
	}
	states = append(states, "RUN")
	js := newJobStateServer(t, states...)
	cli := NewJobSubmitter(nil)
	start := time.Now()
	url, err := cli.waitForInteractiveURL(request.NewAPIClient(js.srv.URL), "tok", "123")
	elapsed := time.Since(start)
	if err != nil {
		t.Fatalf("waitForInteractiveURL: %v", err)
	}
	if url != "http://h:9000" {
		t.Errorf("url = %q, want http://h:9000", url)
	}
	// ~7s expected; under 6s means the backoff never happened (all fast),
	// over 9s means the fast window was skipped (all 2s).
	if elapsed < 6*time.Second || elapsed > 9*time.Second {
		t.Errorf("elapsed = %v, want ~7s (10 fast intervals then a 2s backoff)", elapsed)
	}
	if n := js.pollCount(); n != 13 {
		t.Errorf("polls = %d, want 13 (12 status + 1 url query)", n)
	}
}

// RUN-before-URL polls at the 500ms run cadence: the start-report push
// normally closes the race in well under a second, so a 2s wait here would
// add that much dead time to every interactive submit.
func TestWaitForInteractiveURL_RunRaceUsesFastInterval(t *testing.T) {
	if config.Log == nil {
		config.Log = zap.NewNop()
		t.Cleanup(func() { config.Log = nil })
	}
	// First poll sees RUN with the URL already missing (the state server
	// only exposes a URL once the state flips), second sees it present.
	// We emulate the missing-URL round by answering the first url query
	// with an empty URL: swap in a custom server.
	seen := 0
	emptyURL := true
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		url := ""
		if !emptyURL {
			url = "http://h:9000"
		}
		row := models.QueryJob{JobStatus: "RUN", InteractiveJobURL: url}
		if r.URL.Query().Get("fields") == "interactive_job_url" {
			seen++
			if seen == 1 {
				row.InteractiveJobURL = ""
			}
			emptyURL = row.InteractiveJobURL != ""
		}
		_ = json.NewEncoder(w).Encode(models.JobsResponse{
			Code: 0,
			Data: struct {
				Count      int               `json:"count"`
				Rows       []models.QueryJob `json:"rows"`
				Page_token string            `json:"page_token"`
			}{Rows: []models.QueryJob{row}},
		})
	}))
	t.Cleanup(srv.Close)
	cli := NewJobSubmitter(nil)
	start := time.Now()
	url, err := cli.waitForInteractiveURL(request.NewAPIClient(srv.URL), "tok", "123")
	elapsed := time.Since(start)
	if err != nil {
		t.Fatalf("waitForInteractiveURL: %v", err)
	}
	if url != "http://h:9000" {
		t.Errorf("url = %q, want http://h:9000", url)
	}
	if elapsed > 1200*time.Millisecond {
		t.Errorf("elapsed = %v, want ~0.5s (run cadence, not the old 2s)", elapsed)
	}
}

// A suspended interactive job must not fail the wait the way a terminal
// state does: the allocation is held and `xce jobs resume` may follow, so
// the loop keeps waiting (backed off) and picks up the URL when the job is
// running again.
func TestWaitForInteractiveURL_SuspendedJobKeepsWaiting(t *testing.T) {
	if config.Log == nil {
		config.Log = zap.NewNop()
		t.Cleanup(func() { config.Log = nil })
	}
	js := newJobStateServer(t, "USUSP", "RUN")
	cli := NewJobSubmitter(nil)
	url, err := cli.waitForInteractiveURL(request.NewAPIClient(js.srv.URL), "tok", "123")
	if err != nil {
		t.Fatalf("waitForInteractiveURL: %v", err)
	}
	if url != "http://h:9000" {
		t.Errorf("url = %q, want http://h:9000", url)
	}
}