package-updater/internal/agent/conn/conn_test.go

package conn

import (
	"context"
	"crypto/tls"
	"io"
	"log/slog"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/coder/websocket"
	"github.com/coder/websocket/wsjson"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/shcizo/package-updater/internal/agent/updater"
	"github.com/shcizo/package-updater/internal/pin"
	"github.com/shcizo/package-updater/internal/protocol"
)

type fakeMaster struct {
	srv         *httptest.Server
	fingerprint string
	envs        chan protocol.Envelope
	auth        chan string
	conns       chan *websocket.Conn
	attempts    atomic.Int64
}

func newFakeMaster(t *testing.T) *fakeMaster {
	t.Helper()
	key, err := pin.GenerateKey()
	require.NoError(t, err)
	cert, err := pin.Certificate(key)
	require.NoError(t, err)
	fp, err := pin.Fingerprint(key.Public())
	require.NoError(t, err)

	fm := &fakeMaster{
		fingerprint: fp,
		envs:        make(chan protocol.Envelope, 100),
		auth:        make(chan string, 10),
		conns:       make(chan *websocket.Conn, 10),
	}
	fm.srv = httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fm.auth <- r.Header.Get("Authorization")
		c, err := websocket.Accept(w, r, nil)
		if err != nil {
			return
		}
		c.SetReadLimit(1 << 20)
		fm.conns <- c
		for {
			var env protocol.Envelope
			if err := wsjson.Read(r.Context(), c, &env); err != nil {
				return
			}
			fm.envs <- env
		}
	}))
	fm.srv.Config.ConnState = func(_ net.Conn, state http.ConnState) {
		if state == http.StateNew {
			fm.attempts.Add(1)
		}
	}
	fm.srv.TLS = &tls.Config{Certificates: []tls.Certificate{cert}}
	fm.srv.StartTLS()
	t.Cleanup(fm.srv.Close)
	return fm
}

func (fm *fakeMaster) url() string { return "wss" + strings.TrimPrefix(fm.srv.URL, "https") }

func (fm *fakeMaster) next(t *testing.T) protocol.Envelope {
	t.Helper()
	select {
	case e := <-fm.envs:
		return e
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for a message from the agent")
		return protocol.Envelope{}
	}
}

type fakeSource struct {
	mu       sync.Mutex
	services []protocol.Service
}

func (f *fakeSource) set(s ...protocol.Service) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.services = s
}

func (f *fakeSource) Hello(context.Context) (protocol.Hello, error) {
	return protocol.Hello{Hostname: "node-2"}, nil
}

func (f *fakeSource) Snapshot(context.Context) ([]protocol.Service, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	return append([]protocol.Service(nil), f.services...), nil
}

var (
	svcA = protocol.Service{Kind: protocol.KindContainer, ContainerName: "a", Image: "a:1", Status: "running"}
	svcB = protocol.Service{Kind: protocol.KindContainer, ContainerName: "b", Image: "b:1", Status: "running"}
)

// fakeJobs records what the connection tells it and returns scripted state.
type fakeJobs struct {
	mu       sync.Mutex
	known    []string
	pending  []updater.Finished
	progress []string // output chunks since the last TakeProgress
	submits  map[string]protocol.JobStart
	acks     []string
	changed  chan struct{}
}

func newFakeJobs() *fakeJobs {
	return &fakeJobs{submits: map[string]protocol.JobStart{}, changed: make(chan struct{}, 1)}
}

func (f *fakeJobs) Submit(id string, s protocol.JobStart) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.submits[id] = s
}

func (f *fakeJobs) Ack(id string) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.acks = append(f.acks, id)
	kept := f.pending[:0:0]
	for _, p := range f.pending {
		if p.ID != id {
			kept = append(kept, p)
		}
	}
	f.pending = kept
}

func (f *fakeJobs) Known() []string {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.known
}

func (f *fakeJobs) TakeProgress() []updater.Progress {
	f.mu.Lock()
	defer f.mu.Unlock()
	if len(f.progress) == 0 {
		return nil
	}
	out := updater.Progress{ID: "1", Output: strings.Join(f.progress, "")}
	f.progress = nil
	return []updater.Progress{out}
}

func (f *fakeJobs) Pending() []updater.Finished {
	f.mu.Lock()
	defer f.mu.Unlock()
	return append([]updater.Finished(nil), f.pending...)
}

func (f *fakeJobs) Changed() <-chan struct{} { return f.changed }

func (f *fakeJobs) signal() {
	select {
	case f.changed <- struct{}{}:
	default:
	}
}

func (f *fakeJobs) addProgress(s string) {
	f.mu.Lock()
	f.progress = append(f.progress, s)
	f.mu.Unlock()
	f.signal()
}

func (f *fakeJobs) addResult(id string, r protocol.JobResult) {
	f.mu.Lock()
	f.pending = append(f.pending, updater.Finished{ID: id, Result: r})
	f.mu.Unlock()
	f.signal()
}

func (f *fakeJobs) submitted(id string) (protocol.JobStart, bool) {
	f.mu.Lock()
	defer f.mu.Unlock()
	s, ok := f.submits[id]
	return s, ok
}

func (f *fakeJobs) acked() []string {
	f.mu.Lock()
	defer f.mu.Unlock()
	return append([]string(nil), f.acks...)
}

func startAgent(t *testing.T, fm *fakeMaster, fingerprint string, src Source) *Agent {
	t.Helper()
	return startAgentWith(t, fm, fingerprint, src, newFakeJobs(), 0)
}

func startAgentWith(t *testing.T, fm *fakeMaster, fingerprint string, src Source, jobs Jobs, progressInterval time.Duration) *Agent {
	t.Helper()
	a := New(Config{
		URL: fm.url(), Token: "pua_test", Fingerprint: fingerprint, Version: "v-test",
		MinBackoff: 10 * time.Millisecond, MaxBackoff: 50 * time.Millisecond,
		ProgressInterval: progressInterval,
	}, src, jobs, slog.New(slog.NewTextHandler(io.Discard, nil)))
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() { _ = a.Run(ctx); close(done) }()
	t.Cleanup(func() { cancel(); <-done })
	return a
}

func TestAgent_SendsHelloThenSnapshotWithToken(t *testing.T) {
	fm := newFakeMaster(t)
	src := &fakeSource{}
	src.set(svcA)
	startAgent(t, fm, fm.fingerprint, src)

	assert.Equal(t, "Bearer pua_test", <-fm.auth)
	hello := fm.next(t)
	require.Equal(t, protocol.TypeHello, hello.Type)
	h, err := protocol.Decode[protocol.Hello](hello)
	require.NoError(t, err)
	assert.Equal(t, protocol.Hello{Hostname: "node-2", AgentVersion: "v-test", ProtocolVersion: protocol.Version}, h)

	snap := fm.next(t)
	require.Equal(t, protocol.TypeSnapshot, snap.Type)
	s, err := protocol.Decode[protocol.Snapshot](snap)
	require.NoError(t, err)
	assert.Equal(t, []protocol.Service{svcA}, s.Services)
}

func TestAgent_NotifySendsDiff(t *testing.T) {
	fm := newFakeMaster(t)
	src := &fakeSource{}
	src.set(svcA)
	a := startAgent(t, fm, fm.fingerprint, src)
	fm.next(t) // hello
	fm.next(t) // snapshot

	src.set(svcB)
	a.Notify()

	up := fm.next(t)
	require.Equal(t, protocol.TypeUpsert, up.Type)
	u, err := protocol.Decode[protocol.Upsert](up)
	require.NoError(t, err)
	assert.Equal(t, svcB, u.Service)

	rm := fm.next(t)
	require.Equal(t, protocol.TypeRemove, rm.Type)
	r, err := protocol.Decode[protocol.Remove](rm)
	require.NoError(t, err)
	assert.Equal(t, "container|a", r.Key)
}

func TestAgent_ReconnectsAndResendsSnapshot(t *testing.T) {
	fm := newFakeMaster(t)
	src := &fakeSource{}
	src.set(svcA)
	startAgent(t, fm, fm.fingerprint, src)
	fm.next(t)
	fm.next(t)

	first := <-fm.conns
	first.CloseNow()

	assert.Equal(t, protocol.TypeHello, fm.next(t).Type)
	assert.Equal(t, protocol.TypeSnapshot, fm.next(t).Type)
}

func TestAgent_RefusesMasterWithWrongKey(t *testing.T) {
	fm := newFakeMaster(t)
	other, err := pin.GenerateKey()
	require.NoError(t, err)
	wrong, err := pin.Fingerprint(other.Public())
	require.NoError(t, err)

	startAgent(t, fm, wrong, &fakeSource{})
	require.Eventually(t, func() bool { return fm.attempts.Load() >= 1 },
		2*time.Second, 10*time.Millisecond, "the agent must actually try to connect")
	time.Sleep(200 * time.Millisecond)
	assert.Empty(t, fm.auth, "the token must never be sent to a server with the wrong key")
}

func sendEnv(t *testing.T, c *websocket.Conn, env protocol.Envelope) {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	require.NoError(t, wsjson.Write(ctx, c, env))
}

func TestAgent_HelloCarriesVersionAndKnownJobs(t *testing.T) {
	fm := newFakeMaster(t)
	jobs := newFakeJobs()
	jobs.known = []string{"3", "7"}
	startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond)

	h, err := protocol.Decode[protocol.Hello](fm.next(t))
	require.NoError(t, err)
	assert.Equal(t, protocol.Version, h.ProtocolVersion)
	assert.Equal(t, []string{"3", "7"}, h.KnownJobs)
}

func TestAgent_JobStartSubmitsAndAckAcks(t *testing.T) {
	fm := newFakeMaster(t)
	jobs := newFakeJobs()
	startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond)
	fm.next(t)
	fm.next(t)
	c := <-fm.conns

	start := protocol.JobStart{Kind: protocol.KindSwarm, SwarmClusterID: "c1", SwarmService: "gitea_gitea", Ref: "gitea/gitea:1.27"}
	env, err := protocol.EncodeJob(protocol.TypeJobStart, "42", start)
	require.NoError(t, err)
	sendEnv(t, c, env)
	require.Eventually(t, func() bool { _, ok := jobs.submitted("42"); return ok }, 2*time.Second, 5*time.Millisecond)
	got, _ := jobs.submitted("42")
	assert.Equal(t, start, got)

	ack, err := protocol.EncodeJob(protocol.TypeJobAck, "42", protocol.JobAck{})
	require.NoError(t, err)
	sendEnv(t, c, ack)
	require.Eventually(t, func() bool { return len(jobs.acked()) == 1 }, 2*time.Second, 5*time.Millisecond)
	assert.Equal(t, []string{"42"}, jobs.acked())
}

func TestAgent_SendsPendingResultsAfterSnapshot(t *testing.T) {
	fm := newFakeMaster(t)
	jobs := newFakeJobs()
	jobs.pending = []updater.Finished{
		{ID: "5", Result: protocol.JobResult{Status: protocol.ResultSucceeded, Ref: "a:1@sha256:aa", Output: "ok\n"}},
		{ID: "6", Result: protocol.JobResult{Status: protocol.ResultFailed, Error: "boom"}},
	}
	startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond)

	assert.Equal(t, protocol.TypeHello, fm.next(t).Type)
	assert.Equal(t, protocol.TypeSnapshot, fm.next(t).Type)
	r5 := fm.next(t)
	assert.Equal(t, protocol.TypeJobResult, r5.Type)
	assert.Equal(t, "5", r5.ID)
	res, err := protocol.Decode[protocol.JobResult](r5)
	require.NoError(t, err)
	assert.Equal(t, "a:1@sha256:aa", res.Ref)
	r6 := fm.next(t)
	assert.Equal(t, "6", r6.ID)
}

func TestAgent_ResendsUnackedResultAfterReconnect(t *testing.T) {
	fm := newFakeMaster(t)
	jobs := newFakeJobs()
	jobs.addResult("5", protocol.JobResult{Status: protocol.ResultSucceeded})
	startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond)

	fm.next(t) // hello
	fm.next(t) // snapshot
	assert.Equal(t, "5", fm.next(t).ID)
	(<-fm.conns).CloseNow() // dropped without ack

	assert.Equal(t, protocol.TypeHello, fm.next(t).Type)
	assert.Equal(t, protocol.TypeSnapshot, fm.next(t).Type)
	again := fm.next(t)
	assert.Equal(t, protocol.TypeJobResult, again.Type)
	assert.Equal(t, "5", again.ID)

	c := <-fm.conns
	ack, err := protocol.EncodeJob(protocol.TypeJobAck, "5", protocol.JobAck{})
	require.NoError(t, err)
	sendEnv(t, c, ack)
	require.Eventually(t, func() bool { return len(jobs.acked()) == 1 }, 2*time.Second, 5*time.Millisecond)
	c.CloseNow()

	assert.Equal(t, protocol.TypeHello, fm.next(t).Type)
	assert.Equal(t, protocol.TypeSnapshot, fm.next(t).Type)
	select {
	case e := <-fm.envs:
		t.Fatalf("acked result was sent again: %+v", e)
	case <-time.After(100 * time.Millisecond):
	}
}

func TestAgent_ProgressIsThrottled(t *testing.T) {
	fm := newFakeMaster(t)
	jobs := newFakeJobs()
	startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, 100*time.Millisecond)
	fm.next(t)
	fm.next(t)

	for range 20 {
		jobs.addProgress("x")
		time.Sleep(2500 * time.Microsecond)
	}
	var batches int
	var total strings.Builder
	collect := func(d time.Duration) {
		deadline := time.After(d)
		for {
			select {
			case e := <-fm.envs:
				if e.Type != protocol.TypeJobProgress {
					continue
				}
				p, err := protocol.Decode[protocol.JobProgress](e)
				require.NoError(t, err)
				batches++
				total.WriteString(p.Output)
			case <-deadline:
				return
			}
		}
	}
	collect(100 * time.Millisecond)
	assert.LessOrEqual(t, batches, 2)
	collect(400 * time.Millisecond)
	assert.Equal(t, strings.Repeat("x", 20), total.String(), "the last chunk must be delivered")
}

func TestAgent_IgnoresUnknownType(t *testing.T) {
	fm := newFakeMaster(t)
	jobs := newFakeJobs()
	startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond)
	fm.next(t)
	fm.next(t)
	c := <-fm.conns

	sendEnv(t, c, protocol.Envelope{Type: "future.thing", Payload: []byte(`{"a":1}`)})
	env, err := protocol.EncodeJob(protocol.TypeJobStart, "9", protocol.JobStart{Kind: protocol.KindCompose, Ref: "a:1"})
	require.NoError(t, err)
	sendEnv(t, c, env)
	require.Eventually(t, func() bool { _, ok := jobs.submitted("9"); return ok }, 2*time.Second, 5*time.Millisecond)
	assert.Zero(t, len(fm.conns), "no reconnect expected")
}

func TestAgent_ClosesOnMalformedJobStart(t *testing.T) {
	fm := newFakeMaster(t)
	startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, newFakeJobs(), time.Millisecond)
	fm.next(t)
	fm.next(t)
	c := <-fm.conns

	env, err := protocol.Encode(protocol.TypeJobStart, protocol.JobStart{Ref: "a:1"}) // no ID
	require.NoError(t, err)
	sendEnv(t, c, env)

	select {
	case <-fm.conns:
	case <-time.After(2 * time.Second):
		t.Fatal("agent did not reconnect after a malformed job.start")
	}
}

func TestAgent_InventoryAndProgressConcurrently(t *testing.T) {
	fm := newFakeMaster(t)
	src := &fakeSource{}
	src.set(svcA)
	jobs := newFakeJobs()
	a := startAgentWith(t, fm, fm.fingerprint, src, jobs, time.Millisecond)
	fm.next(t)
	fm.next(t)

	done := make(chan struct{})
	go func() {
		defer close(done)
		for i := range 30 {
			jobs.addProgress("y")
			if i%2 == 0 {
				src.set(svcB)
			} else {
				src.set(svcA)
			}
			a.Notify()
			time.Sleep(time.Millisecond)
		}
	}()
	<-done
	var progress, inventory int
	deadline := time.After(500 * time.Millisecond)
	for {
		select {
		case e := <-fm.envs:
			switch e.Type {
			case protocol.TypeJobProgress:
				progress++
			case protocol.TypeUpsert, protocol.TypeRemove:
				inventory++
			}
		case <-deadline:
			assert.Positive(t, progress)
			assert.Positive(t, inventory)
			return
		}
	}
}