feat: add auto update

server/internal/grpc/registry_test.go

@@ -1,6 +1,7 @@
 package grpc
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -153,3 +154,112 @@ func TestSend_FullChannel(t *testing.T) {
 		t.Error("Send should return false when channel is full")
 	}
 }
+
+// ── Pending file correlations ──────────────────────────────────────────────────
+
+func TestRegisterPending_UnknownAgent(t *testing.T) {
+	r := NewRegistry()
+	ch := r.RegisterPending("ghost", "cmd1")
+	if ch != nil {
+		t.Error("expected nil channel for unknown agent")
+	}
+}
+
+func TestResolvePending_Success(t *testing.T) {
+	r := NewRegistry()
+	r.Register("id1", "h", "a", "ip", "arch", "os")
+
+	ch := r.RegisterPending("id1", "cmd1")
+	if ch == nil {
+		t.Fatal("expected non-nil channel")
+	}
+
+	result := &agentv1.FileResult{CommandId: "cmd1", Success: true, Content: []byte("data")}
+	r.ResolvePending("id1", "cmd1", result)
+
+	select {
+	case got := <-ch:
+		if got.CommandId != "cmd1" || !got.Success {
+			t.Errorf("unexpected result: %+v", got)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("timed out waiting for resolve")
+	}
+}
+
+func TestResolvePending_UnknownAgent(t *testing.T) {
+	r := NewRegistry()
+	// must not panic
+	r.ResolvePending("ghost", "cmd1", &agentv1.FileResult{})
+}
+
+func TestResolvePending_UnknownCmd(t *testing.T) {
+	r := NewRegistry()
+	r.Register("id1", "h", "a", "ip", "arch", "os")
+	// must not panic
+	r.ResolvePending("id1", "nonexistent", &agentv1.FileResult{})
+}
+
+func TestCancelPending(t *testing.T) {
+	r := NewRegistry()
+	r.Register("id1", "h", "a", "ip", "arch", "os")
+
+	r.RegisterPending("id1", "cmd1")
+	r.CancelPending("id1", "cmd1")
+
+	// After cancel, resolving should be a no-op (not panic)
+	r.ResolvePending("id1", "cmd1", &agentv1.FileResult{})
+}
+
+func TestCancelPending_UnknownAgent(t *testing.T) {
+	r := NewRegistry()
+	// must not panic
+	r.CancelPending("ghost", "cmd1")
+}
+
+func TestSendAndWaitCtx_AgentNotConnected(t *testing.T) {
+	r := NewRegistry()
+	ctx := context.Background()
+	_, err := r.SendAndWaitCtx(ctx, "ghost", &agentv1.ServerMessage{}, "cmd1")
+	if err == nil || err.Error() != "agent not connected" {
+		t.Errorf("expected 'agent not connected', got %v", err)
+	}
+}
+
+func TestSendAndWaitCtx_Timeout(t *testing.T) {
+	r := NewRegistry()
+	r.Register("id1", "h", "a", "ip", "arch", "os")
+
+	// Use an already-cancelled context to force immediate timeout.
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // cancel immediately
+
+	_, err := r.SendAndWaitCtx(ctx, "id1", &agentv1.ServerMessage{}, "cmd-timeout")
+	if err == nil {
+		t.Error("expected timeout or not-connected error")
+	}
+}
+
+func TestSendAndWaitCtx_Success(t *testing.T) {
+	r := NewRegistry()
+	r.Register("id1", "h", "a", "ip", "arch", "os")
+
+	cmdID := "cmd-success"
+	expected := &agentv1.FileResult{CommandId: cmdID, Success: true, Content: []byte("hello")}
+
+	// Simulate the agent responding after the send.
+	go func() {
+		// Wait briefly for RegisterPending + Send to happen.
+		time.Sleep(10 * time.Millisecond)
+		r.ResolvePending("id1", cmdID, expected)
+	}()
+
+	ctx := context.Background()
+	result, err := r.SendAndWaitCtx(ctx, "id1", &agentv1.ServerMessage{}, cmdID)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if result.CommandId != cmdID || !result.Success {
+		t.Errorf("unexpected result: %+v", result)
+	}
+}