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/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"}
)

func startAgent(t *testing.T, fm *fakeMaster, fingerprint string, src Source) *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,
	}, src, 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"}, 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")
}