oto/services/core/internal/runnersocket/server_test.go

package runnersocket

import (
	"net"
	"testing"
	"time"

	protoSocket "git.toki-labs.com/toki/proto-socket/go"
	"github.com/toki/oto/services/core/internal/cicdstate"
	"github.com/toki/oto/services/core/internal/runnerregistry"
	otopb "github.com/toki/oto/services/core/oto"
)

func testServer(registry *runnerregistry.Registry) *Server {
	return New("127.0.0.1:0", registry, cicdstate.NewStore())
}

// pipeClient creates a TcpClient backed by one end of a net.Pipe.
// heartbeat is disabled (intervalSec=0) to avoid background timer interference.
// Cleanup closes both ends.
func pipeClient(t *testing.T) (*protoSocket.TcpClient, net.Conn) {
	t.Helper()
	a, b := net.Pipe()
	client := protoSocket.NewTcpClient(a, 0, 0, parserMap())
	t.Cleanup(func() { _ = a.Close(); _ = b.Close() })
	return client, b
}

// runnerParserMap returns a parser map for the runner (client) side of a
// proto-socket connection. It includes server-to-runner frames so tests can
// verify dispatch and cancel delivery.
func runnerParserMap() protoSocket.ParserMap {
	return protoSocket.ParserMap{
		protoSocket.TypeNameOf(&otopb.RunRequest{}):       parserFor(func() *otopb.RunRequest { return &otopb.RunRequest{} }),
		protoSocket.TypeNameOf(&otopb.CancelRunRequest{}): parserFor(func() *otopb.CancelRunRequest { return &otopb.CancelRunRequest{} }),
	}
}

// pipeClientPair creates two TcpClients backed by the two ends of a net.Pipe.
// serverSide uses parserMap(); runnerSide uses runnerParserMap() so tests can
// capture and inspect RunRequest frames sent by DispatchQueuedJobs.
// Cleanup closes both connections.
func pipeClientPair(t *testing.T) (serverSide, runnerSide *protoSocket.TcpClient) {
	t.Helper()
	a, b := net.Pipe()
	serverSide = protoSocket.NewTcpClient(a, 0, 0, parserMap())
	runnerSide = protoSocket.NewTcpClient(b, 0, 0, runnerParserMap())
	t.Cleanup(func() { _ = a.Close(); _ = b.Close() })
	return serverSide, runnerSide
}

// TestServerSetClientReplacesDuplicate verifies that registering a second client
// for the same runner ID closes the old one and stores the new one.
func TestServerReplacesDuplicateRunnerClient(t *testing.T) {
	s := testServer(runnerregistry.New())

	client1, _ := pipeClient(t)
	client2, _ := pipeClient(t)

	s.setClient("runner-A", client1)
	s.setClient("runner-A", client2)

	// Allow Close() to propagate through connCloseOnce.
	time.Sleep(20 * time.Millisecond)

	if client1.IsAlive() {
		t.Error("old client must be closed after replacement, but it is still alive")
	}
	if got := s.clientFor("runner-A"); got != client2 {
		t.Errorf("clientFor must return new client after replacement; got %v", got)
	}
}

// TestServerIgnoresStaleClientDisconnect verifies that removeClient ignores
// a client that is not the current entry for the given runner ID.
func TestServerIgnoresStaleClientDisconnect(t *testing.T) {
	s := testServer(runnerregistry.New())

	client1, _ := pipeClient(t)
	client2, _ := pipeClient(t)

	s.setClient("runner-B", client1)

	// client2 is a different object – simulates a stale disconnect listener.
	removed := s.removeClient("runner-B", client2)
	if removed {
		t.Error("removeClient must return false for a client that is not the current one")
	}
	if got := s.clientFor("runner-B"); got != client1 {
		t.Errorf("current client must still be client1 after stale removeClient; got %v", got)
	}
}

// TestServerStaleClientDisconnectDoesNotDisconnectRegistry verifies the complete
// stale-disconnect protection chain:
//
//  1. setClient replaces the old client.
//  2. removeClient returns false for the stale old client.
//  3. Because removeClient returns false, registry.Disconnect is NOT called,
//     so the runner remains in a healthy (non-disconnected) state.
func TestServerStaleClientDisconnectDoesNotDisconnectRegistry(t *testing.T) {
	registry := runnerregistry.New()
	s := testServer(registry)

	resp := registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "tok",
		RunnerId:        "runner-C",
		ProtocolVersion: "oto.runner.v1",
		Capability:      &otopb.RunnerCapability{Name: "oto-runner", Version: "1.0.0"},
	})
	if !resp.GetAccepted() {
		t.Fatalf("registry.Register rejected: %s", resp.GetRejectReason())
	}
	registry.Heartbeat("runner-C", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)

	client1, _ := pipeClient(t)
	client2, _ := pipeClient(t)

	s.setClient("runner-C", client1)
	// Duplicate connection: client2 replaces client1.
	s.setClient("runner-C", client2)

	// Simulate stale client1's disconnect listener.
	removed := s.removeClient("runner-C", client1)
	if removed {
		t.Error("stale removeClient should return false")
	}

	// Because removeClient returned false the disconnect listener guard skips
	// registry.Disconnect. The runner must NOT be in disconnected state.
	runner, ok := registry.Snapshot("runner-C")
	if !ok {
		t.Fatal("runner-C not found in registry")
	}
	if runner.Status == runnerregistry.StatusDisconnected {
		t.Errorf("runner status must not be %q after stale disconnect; got %q",
			runnerregistry.StatusDisconnected, runner.Status)
	}
	if got := s.clientFor("runner-C"); got != client2 {
		t.Errorf("current client must still be client2; got %v", got)
	}
}

// TestServerDuplicateReplaceWithListenerDoesNotDeadlock verifies that replacing
// a client that carries a production-style disconnect listener does not deadlock.
//
// Before the fix, setClient held s.mu while calling old.Close(), which fired
// the listener synchronously. The listener called removeClient → tried to
// acquire s.mu again → deadlock.  After the fix, old.Close() is called after
// s.mu is released, so the re-entry is safe.
func TestServerDuplicateReplaceWithListenerDoesNotDeadlock(t *testing.T) {
	registry := runnerregistry.New()
	s := testServer(registry)

	resp := registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "tok",
		RunnerId:        "runner-E",
		ProtocolVersion: "oto.runner.v1",
		Capability:      &otopb.RunnerCapability{Name: "oto-runner", Version: "1.0.0"},
	})
	if !resp.GetAccepted() {
		t.Fatalf("registry.Register rejected: %s", resp.GetRejectReason())
	}

	client1, _ := pipeClient(t)
	client2, _ := pipeClient(t)

	// Mirror the production disconnect listener from attachClientHandlers.
	// This is the re-entry path: removeClient acquires s.mu, which would
	// deadlock if setClient still held s.mu during old.Close().
	client1.AddDisconnectListener(func(c *protoSocket.TcpClient) {
		if s.removeClient("runner-E", c) {
			s.registry.Disconnect("runner-E")
		}
	})

	s.setClient("runner-E", client1)

	// This must complete without hanging. If the lock-order bug is present the
	// goroutine stalls here and the test times out via go test -timeout.
	s.setClient("runner-E", client2)

	// Allow Close() and the disconnect listener to propagate.
	time.Sleep(20 * time.Millisecond)

	if client1.IsAlive() {
		t.Error("old client must be closed after replacement")
	}
	if got := s.clientFor("runner-E"); got != client2 {
		t.Errorf("clientFor must return new client after replacement; got %v", got)
	}
	// The stale listener must not have disconnected the runner.
	runner, ok := registry.Snapshot("runner-E")
	if !ok {
		t.Fatal("runner-E not found in registry after replacement")
	}
	if runner.Status == runnerregistry.StatusDisconnected {
		t.Errorf("runner must not be disconnected after stale listener fires; got %q", runner.Status)
	}
}

// testServerWithStore creates a Server sharing the given registry and returns
// the associated Store for state inspection in tests.
func testServerWithStore(registry *runnerregistry.Registry) (*Server, *cicdstate.Store) {
	store := cicdstate.NewStore()
	return New("127.0.0.1:0", registry, store), store
}

// registerRunner adds a runner to the registry with default test fields.
func registerRunner(t *testing.T, registry *runnerregistry.Registry, runnerID string) {
	t.Helper()
	resp := registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "tok",
		RunnerId:        runnerID,
		ProtocolVersion: "oto.runner.v1",
		Capability:      &otopb.RunnerCapability{Name: "oto-runner", Version: "1.0.0"},
	})
	if !resp.GetAccepted() {
		t.Fatalf("registerRunner %s: rejected: %s", runnerID, resp.GetRejectReason())
	}
}

// createQueuedJob adds a job with a valid inline YAML to the store.
func createQueuedJob(t *testing.T, store *cicdstate.Store, jobID string) {
	t.Helper()
	if _, err := store.CreateJob(jobID, jobID, &cicdstate.RunInput{PipelineYAML: "pipeline: test"}); err != nil {
		t.Fatalf("CreateJob %s: %v", jobID, err)
	}
}

// setupRunningExecution creates a queued job, registers the runner, creates an
// execution, and transitions both job and execution to running. Returns execID.
func setupRunningExecution(t *testing.T, store *cicdstate.Store, registry *runnerregistry.Registry, jobID, runnerID string) string {
	t.Helper()
	createQueuedJob(t, store, jobID)
	registerRunner(t, registry, runnerID)
	execID := store.NextExecutionID()
	if _, err := store.CreateExecution(jobID, execID); err != nil {
		t.Fatalf("CreateExecution: %v", err)
	}
	_ = store.SetExecutionRunnerID(execID, runnerID)
	if err := store.TransitionJob(jobID, cicdstate.StateRunning); err != nil {
		t.Fatalf("TransitionJob running: %v", err)
	}
	if err := store.TransitionExecution(execID, cicdstate.StateRunning); err != nil {
		t.Fatalf("TransitionExecution running: %v", err)
	}
	return execID
}

// TestServerDispatchQueuedJobOverSocket verifies that DispatchQueuedJobs
// transitions the queued job and its execution to running over the socket
// transport, assigns the dispatched runner ID to the execution, and sends a
// RunRequest frame whose fields match the store state.
func TestServerDispatchQueuedJobOverSocket(t *testing.T) {
	registry := runnerregistry.New()
	s, store := testServerWithStore(registry)
	registerRunner(t, registry, "runner-dispatch")

	jobID := "job-dispatch"
	if _, err := store.CreateJob(jobID, jobID, &cicdstate.RunInput{
		PipelineYAML: "pipeline: dispatch-test",
		Variables:    map[string]string{"ENV": "ci"},
		CommandTypes: []string{"flutter"},
	}); err != nil {
		t.Fatalf("CreateJob: %v", err)
	}

	// Use a runner-side TcpClient so we can parse the RunRequest frame.
	serverClient, runnerClient := pipeClientPair(t)

	received := make(chan *otopb.RunRequest, 1)
	protoSocket.AddListenerTyped[*otopb.RunRequest](&runnerClient.Communicator, func(req *otopb.RunRequest) {
		select {
		case received <- req:
		default:
		}
	})

	s.setClient("runner-dispatch", serverClient)
	s.DispatchQueuedJobs()

	var rr *otopb.RunRequest
	select {
	case rr = <-received:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for RunRequest")
	}

	job, err := store.GetJob(jobID)
	if err != nil {
		t.Fatalf("GetJob: %v", err)
	}
	if job.State != cicdstate.StateRunning {
		t.Errorf("job state: want %s, got %s", cicdstate.StateRunning, job.State)
	}
	if job.ExecutionID == "" {
		t.Fatal("no execution created for dispatched job")
	}
	exec, err := store.GetExecution(job.ExecutionID)
	if err != nil {
		t.Fatalf("GetExecution: %v", err)
	}
	if exec.State != cicdstate.StateRunning {
		t.Errorf("execution state: want %s, got %s", cicdstate.StateRunning, exec.State)
	}
	if exec.RunnerID != "runner-dispatch" {
		t.Errorf("execution runnerID: want runner-dispatch, got %s", exec.RunnerID)
	}

	// Verify RunRequest payload fields match the store state.
	if rr.RunnerId != "runner-dispatch" {
		t.Errorf("RunRequest.RunnerId: want runner-dispatch, got %s", rr.RunnerId)
	}
	if rr.JobId != jobID {
		t.Errorf("RunRequest.JobId: want %s, got %s", jobID, rr.JobId)
	}
	if rr.ExecutionId != job.ExecutionID {
		t.Errorf("RunRequest.ExecutionId: want %s, got %s", job.ExecutionID, rr.ExecutionId)
	}
	if rr.PipelineYaml != "pipeline: dispatch-test" {
		t.Errorf("RunRequest.PipelineYaml: want %q, got %q", "pipeline: dispatch-test", rr.PipelineYaml)
	}
	if len(rr.Variables) == 0 || rr.Variables["ENV"] != "ci" {
		t.Errorf("RunRequest.Variables: want {ENV:ci}, got %v", rr.Variables)
	}
	if len(rr.CommandTypes) != 1 || rr.CommandTypes[0] != "flutter" {
		t.Errorf("RunRequest.CommandTypes: want [flutter], got %v", rr.CommandTypes)
	}
}

// TestServerCancelExecutionSendsSocketCancelToRunner verifies the server-side
// cancel action uses the runner-facing socket as a best-effort signal channel.
func TestServerCancelExecutionSendsSocketCancelToRunner(t *testing.T) {
	s := testServer(runnerregistry.New())
	serverClient, runnerClient := pipeClientPair(t)

	received := make(chan *otopb.CancelRunRequest, 1)
	protoSocket.AddListenerTyped[*otopb.CancelRunRequest](&runnerClient.Communicator, func(req *otopb.CancelRunRequest) {
		select {
		case received <- req:
		default:
		}
	})

	s.setClient("runner-cancel", serverClient)

	if err := s.CancelExecution("runner-cancel", "exec-cancel", "operator request"); err != nil {
		t.Fatalf("CancelExecution: %v", err)
	}

	var req *otopb.CancelRunRequest
	select {
	case req = <-received:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for CancelRunRequest")
	}
	if req.RunnerId != "runner-cancel" {
		t.Errorf("CancelRunRequest.RunnerId: want runner-cancel, got %s", req.RunnerId)
	}
	if req.ExecutionId != "exec-cancel" {
		t.Errorf("CancelRunRequest.ExecutionId: want exec-cancel, got %s", req.ExecutionId)
	}
	if req.Reason != "operator request" {
		t.Errorf("CancelRunRequest.Reason: want operator request, got %s", req.Reason)
	}

	if err := s.CancelExecution("runner-missing", "exec-cancel", "operator request"); err == nil {
		t.Fatal("expected missing connected runner to return an error")
	}
}

// TestServerDispatchClaimsOnlyOneJobPerRunner verifies that a single
// DispatchQueuedJobs call dispatches at most one job per runner, leaving any
// remaining queued jobs untouched.
func TestServerDispatchClaimsOnlyOneJobPerRunner(t *testing.T) {
	registry := runnerregistry.New()
	s, store := testServerWithStore(registry)
	registerRunner(t, registry, "runner-claim")
	createQueuedJob(t, store, "job-claim-1")
	createQueuedJob(t, store, "job-claim-2")

	client, remote := pipeClient(t)
	// Drain the remote end so Send does not block.
	go func() {
		buf := make([]byte, 4096)
		for {
			if _, err := remote.Read(buf); err != nil {
				return
			}
		}
	}()

	s.setClient("runner-claim", client)
	s.DispatchQueuedJobs()

	job1, err := store.GetJob("job-claim-1")
	if err != nil {
		t.Fatalf("GetJob job-claim-1: %v", err)
	}
	job2, err := store.GetJob("job-claim-2")
	if err != nil {
		t.Fatalf("GetJob job-claim-2: %v", err)
	}

	running, queued := 0, 0
	for _, j := range []*cicdstate.Job{job1, job2} {
		switch j.State {
		case cicdstate.StateRunning:
			running++
		case cicdstate.StateQueued:
			queued++
		}
	}
	if running != 1 {
		t.Errorf("want exactly 1 running job after one dispatch pass, got %d", running)
	}
	if queued != 1 {
		t.Errorf("want exactly 1 queued job after one dispatch pass, got %d", queued)
	}
}

// TestServerFailsExecutionWhenRunRequestCannotBeSent verifies that a send
// error over the socket transitions the claimed execution and job to failed
// rather than leaving them stuck in running.
func TestServerFailsExecutionWhenRunRequestCannotBeSent(t *testing.T) {
	registry := runnerregistry.New()
	s, store := testServerWithStore(registry)
	registerRunner(t, registry, "runner-sendfail")
	createQueuedJob(t, store, "job-sendfail")

	client, remote := pipeClient(t)
	// Close the remote end before dispatch so any write to the client fails.
	_ = remote.Close()

	s.setClient("runner-sendfail", client)
	s.DispatchQueuedJobs()

	job, err := store.GetJob("job-sendfail")
	if err != nil {
		t.Fatalf("GetJob: %v", err)
	}
	if job.State != cicdstate.StateFailed {
		t.Errorf("job state: want %s, got %s", cicdstate.StateFailed, job.State)
	}
	if job.ExecutionID == "" {
		t.Fatal("expected execution to be created before send failure")
	}
	exec, err := store.GetExecution(job.ExecutionID)
	if err != nil {
		t.Fatalf("GetExecution: %v", err)
	}
	if exec.State != cicdstate.StateFailed {
		t.Errorf("execution state: want %s, got %s", cicdstate.StateFailed, exec.State)
	}
}

// TestServerReportsRunResultOverSocket verifies that handleExecutionReport
// accepts success and failure terminal reports from the owning runner and
// transitions execution and job to the correct terminal state.
func TestServerReportsRunResultOverSocket(t *testing.T) {
	cases := []struct {
		name      string
		success   bool
		wantState string
	}{
		{"success", true, cicdstate.StateSucceeded},
		{"failure", false, cicdstate.StateFailed},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			registry := runnerregistry.New()
			s, store := testServerWithStore(registry)
			jobID := "job-report-" + tc.name
			runnerID := "runner-report-" + tc.name
			execID := setupRunningExecution(t, store, registry, jobID, runnerID)

			resp := s.handleExecutionReport(&otopb.ExecutionReportRequest{
				RunnerId:    runnerID,
				JobId:       jobID,
				ExecutionId: execID,
				Success:     tc.success,
			})

			if !resp.GetAccepted() {
				t.Fatalf("report rejected: %s", resp.GetErrorMessage())
			}
			exec, err := store.GetExecution(execID)
			if err != nil {
				t.Fatalf("GetExecution: %v", err)
			}
			if exec.State != tc.wantState {
				t.Errorf("execution state: want %s, got %s", tc.wantState, exec.State)
			}
			job, err := store.GetJob(jobID)
			if err != nil {
				t.Fatalf("GetJob: %v", err)
			}
			if job.State != tc.wantState {
				t.Errorf("job state: want %s, got %s", tc.wantState, job.State)
			}
		})
	}
}

// TestServerRejectsTerminalReportForInactiveExecution verifies the active
// execution gate: terminal reports on queued, succeeded, or cancelled
// executions are rejected, and reports from the wrong runner are rejected.
func TestServerRejectsTerminalReportForInactiveExecution(t *testing.T) {
	t.Run("queued_execution", func(t *testing.T) {
		registry := runnerregistry.New()
		s, store := testServerWithStore(registry)
		registerRunner(t, registry, "runner-guard")
		createQueuedJob(t, store, "job-guard-queued")
		execID := store.NextExecutionID()
		if _, err := store.CreateExecution("job-guard-queued", execID); err != nil {
			t.Fatalf("CreateExecution: %v", err)
		}
		_ = store.SetExecutionRunnerID(execID, "runner-guard")
		// Execution remains in queued state; terminal report must be rejected.
		resp := s.handleExecutionReport(&otopb.ExecutionReportRequest{
			RunnerId: "runner-guard", JobId: "job-guard-queued", ExecutionId: execID, Success: true,
		})
		if resp.GetAccepted() {
			t.Error("expected rejection for queued execution, got accepted")
		}
	})

	t.Run("succeeded_execution", func(t *testing.T) {
		registry := runnerregistry.New()
		s, store := testServerWithStore(registry)
		execID := setupRunningExecution(t, store, registry, "job-guard-succ", "runner-guard-succ")
		_ = store.TransitionExecution(execID, cicdstate.StateSucceeded)
		_ = store.TransitionJob("job-guard-succ", cicdstate.StateSucceeded)

		resp := s.handleExecutionReport(&otopb.ExecutionReportRequest{
			RunnerId: "runner-guard-succ", JobId: "job-guard-succ", ExecutionId: execID, Success: true,
		})
		if resp.GetAccepted() {
			t.Error("expected rejection for succeeded execution, got accepted")
		}
	})

	t.Run("cancelled_execution", func(t *testing.T) {
		registry := runnerregistry.New()
		s, store := testServerWithStore(registry)
		execID := setupRunningExecution(t, store, registry, "job-guard-cancel", "runner-guard-cancel")
		if err := store.CancelJobExecution("job-guard-cancel", execID); err != nil {
			t.Fatalf("CancelJobExecution: %v", err)
		}
		resp := s.handleExecutionReport(&otopb.ExecutionReportRequest{
			RunnerId: "runner-guard-cancel", JobId: "job-guard-cancel", ExecutionId: execID, Success: true,
		})
		if resp.GetAccepted() {
			t.Error("expected rejection for cancelled execution, got accepted")
		}
	})

	t.Run("wrong_runner", func(t *testing.T) {
		registry := runnerregistry.New()
		s, store := testServerWithStore(registry)
		execID := setupRunningExecution(t, store, registry, "job-guard-runner", "runner-owner")
		registerRunner(t, registry, "runner-other")

		resp := s.handleExecutionReport(&otopb.ExecutionReportRequest{
			RunnerId: "runner-other", JobId: "job-guard-runner", ExecutionId: execID, Success: true,
		})
		if resp.GetAccepted() {
			t.Error("expected rejection for wrong runner, got accepted")
		}
		if resp.GetErrorMessage() == "" {
			t.Error("expected non-empty error message for wrong runner")
		}
	})
}

// TestServerDispatchInvalidRunInputFails verifies that DispatchQueuedJobs
// transitions a job with an invalid RunInput directly to failed without creating
// an execution. Two invalid fixture variants are covered: both fields empty and
// both fields set simultaneously.
func TestServerDispatchInvalidRunInputFails(t *testing.T) {
	cases := []struct {
		name  string
		input *cicdstate.RunInput
	}{
		{
			"both_empty",
			&cicdstate.RunInput{PipelineYAML: "", PipelineYAMLPath: ""},
		},
		{
			"both_set",
			&cicdstate.RunInput{PipelineYAML: "pipeline: test", PipelineYAMLPath: "/some/path.yaml"},
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			registry := runnerregistry.New()
			s, store := testServerWithStore(registry)
			runnerID := "runner-invalid-" + tc.name
			jobID := "job-invalid-" + tc.name
			registerRunner(t, registry, runnerID)

			if _, err := store.CreateJob(jobID, jobID, tc.input); err != nil {
				t.Fatalf("CreateJob: %v", err)
			}

			// No drain needed: validateRunInput fails before client.Send is called.
			client, _ := pipeClient(t)
			s.setClient(runnerID, client)
			s.DispatchQueuedJobs()

			job, err := store.GetJob(jobID)
			if err != nil {
				t.Fatalf("GetJob: %v", err)
			}
			if job.State != cicdstate.StateFailed {
				t.Errorf("job state: want %s, got %s", cicdstate.StateFailed, job.State)
			}
			if job.ExecutionID != "" {
				t.Errorf("no execution should be created for invalid run input; got execID %s", job.ExecutionID)
			}
		})
	}
}

// TestServerReportMessageSavedToLog verifies that a non-empty Message in an
// ExecutionReportRequest is appended to the execution log after the report is
// accepted.
func TestServerReportMessageSavedToLog(t *testing.T) {
	cases := []struct {
		name    string
		success bool
		message string
	}{
		{"success_with_message", true, "build succeeded: artifacts uploaded"},
		{"failure_with_message", false, "build failed: compilation error"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			registry := runnerregistry.New()
			s, store := testServerWithStore(registry)
			jobID := "job-log-" + tc.name
			runnerID := "runner-log-" + tc.name
			execID := setupRunningExecution(t, store, registry, jobID, runnerID)

			resp := s.handleExecutionReport(&otopb.ExecutionReportRequest{
				RunnerId:    runnerID,
				JobId:       jobID,
				ExecutionId: execID,
				Success:     tc.success,
				Message:     tc.message,
			})
			if !resp.GetAccepted() {
				t.Fatalf("report rejected: %s", resp.GetErrorMessage())
			}

			logs, err := store.GetLogs(execID)
			if err != nil {
				t.Fatalf("GetLogs: %v", err)
			}
			found := false
			for _, entry := range logs {
				if entry.Line == tc.message {
					found = true
					break
				}
			}
			if !found {
				t.Errorf("expected execution log to contain %q; got entries: %v", tc.message, logs)
			}
		})
	}
}

// TestServerHeartbeatUsesRegisteredRunnerID verifies that connectedRunnerIDs
// reflects only currently active clients and that clientFor returns nil after
// the current client is removed.
func TestServerHeartbeatUsesRegisteredRunnerID(t *testing.T) {
	s := testServer(runnerregistry.New())

	client1, _ := pipeClient(t)

	if ids := s.connectedRunnerIDs(); len(ids) != 0 {
		t.Errorf("expected no connected runners before any setClient; got %v", ids)
	}

	s.setClient("runner-D", client1)

	ids := s.connectedRunnerIDs()
	if len(ids) != 1 || ids[0] != "runner-D" {
		t.Errorf("expected [runner-D] after setClient; got %v", ids)
	}

	// Current client removed → runner no longer in connected list.
	s.removeClient("runner-D", client1)

	if ids := s.connectedRunnerIDs(); len(ids) != 0 {
		t.Errorf("expected no connected runners after removeClient; got %v", ids)
	}
	if got := s.clientFor("runner-D"); got != nil {
		t.Errorf("expected nil from clientFor after removeClient; got %v", got)
	}
}