From a278bf77d68b06a5ea7cb717c59da90c20f12063 Mon Sep 17 00:00:00 2001 From: Scott Funkenhauser Date: Wed, 26 Aug 2026 15:05:23 -0400 Subject: [PATCH] Add functional tests for worker API --- .../controlapi/functionaltest/worker_test.go | 192 +++++++++++++++++- 1 file changed, 186 insertions(+), 6 deletions(-) diff --git a/cmd/ateapi/internal/controlapi/functionaltest/worker_test.go b/cmd/ateapi/internal/controlapi/functionaltest/worker_test.go index f36cd14334..b403846276 100644 --- a/cmd/ateapi/internal/controlapi/functionaltest/worker_test.go +++ b/cmd/ateapi/internal/controlapi/functionaltest/worker_test.go @@ -16,21 +16,117 @@ package functionaltest import ( "context" + "fmt" "testing" + "time" "github.com/agent-substrate/substrate/pkg/proto/ateapipb" "github.com/google/go-cmp/cmp" "google.golang.org/grpc/codes" "google.golang.org/protobuf/testing/protocmp" + "k8s.io/apimachinery/pkg/util/wait" ) -// TestListWorkers tests that workers mirrored to the store are listed. +// Workers are named after the pod UID they were derived from. +const testWorkerName = "5f2c1a90-7b34-4e6d-8a11-0c3e9d5b7f42" + +// newTestWorker returns a Worker in the shape CreateWorker accepts: named after +// its pod UID, with the pod coordinates filled in and no status, which is +// output-only. +// +// The Worker stands on its own, with no pod behind it. That is enough for the +// CRUD API — only the atelet dialer resolves a Worker to a pod, and nothing in +// ate-api reconciles the two now that the syncer runs in ate-controller. Tests +// that need both use createWorkerPod. +func newTestWorker(ns string) *ateapipb.Worker { + return &ateapipb.Worker{ + Metadata: &ateapipb.ResourceMetadata{Name: testWorkerName}, + WorkerNamespace: ns, + WorkerPool: "pool1", + WorkerPod: "worker-api-1", + WorkerPodUid: testWorkerName, + NodeName: "node1", + Ip: "10.1.2.3", + SandboxClass: "gvisor", + Capacity: &ateapipb.WorkerCapacity{CpuMilli: 2000, MemoryBytes: 4 << 30}, + } +} + +// registerWorker registers newTestWorker over the API and returns it as stored. +func registerWorker(t *testing.T, tc *testContext, ns string) *ateapipb.Worker { + t.Helper() + created, err := tc.client.CreateWorker(context.Background(), &ateapipb.CreateWorkerRequest{Worker: newTestWorker(ns)}) + if err != nil { + t.Fatalf("CreateWorker failed: %v", err) + } + return created +} + +func workerRef(name string) *ateapipb.ObjectRef { + return &ateapipb.ObjectRef{Name: name} +} + +// waitForWorkerState waits for a state committed by an RPC to reach the worker +// cache, which follows the store through a PostgreSQL watch and so lands +// shortly after the call returns rather than with it. +func waitForWorkerState(t *testing.T, tc *testContext, name string, want ateapipb.WorkerState) { + t.Helper() + err := wait.PollUntilContextTimeout(context.Background(), 10*time.Millisecond, 5*time.Second, true, func(context.Context) (bool, error) { + worker, err := tc.workerCache.Worker(name) + if err != nil { + return false, nil + } + return worker.GetStatus().GetState() == want, nil + }) + if err != nil { + t.Fatalf("worker %s did not reach state %v in the worker cache: %v", name, want, err) + } +} + +func containsWorker(workers []*ateapipb.Worker, name string) bool { + for _, w := range workers { + if w.GetMetadata().GetName() == name { + return true + } + } + return false +} + +// TestCreateAndGetWorker registers a Worker and reads it back, checking that +// the server assigns the metadata and the ACTIVE status a new Worker starts in. +func TestCreateAndGetWorker(t *testing.T) { + ns := namespaceForTest("ns-worker-create-get") + tc := setupTest(t, ns) + defer tc.cleanup() + + created := registerWorker(t, tc, ns) + + want := newTestWorker(ns) + want.Metadata.Version = 1 + want.Status = &ateapipb.WorkerStatus{State: ateapipb.WorkerState_WORKER_STATE_ACTIVE} + if diff := cmp.Diff(want, created, protocmp.Transform(), ignoreServerMetadata); diff != "" { + t.Errorf("CreateWorker response mismatch (-want +got):\n%s", diff) + } + if created.GetMetadata().GetUid() == "" { + t.Errorf("CreateWorker returned no uid, which every later guard needs") + } + + got, err := tc.client.GetWorker(context.Background(), &ateapipb.GetWorkerRequest{Worker: workerRef(testWorkerName)}) + if err != nil { + t.Fatalf("GetWorker failed: %v", err) + } + if diff := cmp.Diff(created, got, protocmp.Transform()); diff != "" { + t.Errorf("GetWorker returned something other than what CreateWorker stored (-created +got):\n%s", diff) + } +} + +// TestListWorkers tests that registered workers are listed. // Workflow: -// 1. Creates a mock WorkerPool in Kubernetes. -// 2. Creates a mock worker Pod in Kubernetes belonging to that pool. -// 3. Waits for the background WorkerPoolSyncer to mirror it to the store. -// 4. Calls ListWorkers RPC. -// 5. Verifies that the worker appears in the response. +// 1. Creates a mock WorkerPool in Kubernetes. +// 2. Creates a mock worker Pod in Kubernetes belonging to that pool, and +// registers the Worker the syncer would derive from it. +// 3. Calls ListWorkers RPC. +// 4. Verifies that the worker appears in the response. func TestListWorkers(t *testing.T) { ns := namespaceForTest("ns-list-workers") tc := setupTest(t, ns) @@ -76,6 +172,90 @@ func TestListWorkers(t *testing.T) { } } +// TestUpdateWorker changes the two mutable fields. An update replaces the whole +// Worker, so the request is the observed one with those fields changed: its +// metadata carries the uid and version guards, and everything else has to be +// sent back as-is to avoid being cleared. +func TestUpdateWorker(t *testing.T) { + ns := namespaceForTest("ns-worker-update") + tc := setupTest(t, ns) + defer tc.cleanup() + + worker := registerWorker(t, tc, ns) + worker.SandboxClass = "runc" + worker.Labels = map[string]string{"tier": "batch"} + + updated, err := tc.client.UpdateWorker(context.Background(), &ateapipb.UpdateWorkerRequest{Worker: worker}) + if err != nil { + t.Fatalf("UpdateWorker failed: %v", err) + } + + want := newTestWorker(ns) + want.Metadata.Version = 2 + want.SandboxClass = "runc" + want.Labels = map[string]string{"tier": "batch"} + want.Status = &ateapipb.WorkerStatus{State: ateapipb.WorkerState_WORKER_STATE_ACTIVE} + if diff := cmp.Diff(want, updated, protocmp.Transform(), ignoreServerMetadata); diff != "" { + t.Errorf("UpdateWorker response mismatch (-want +got):\n%s", diff) + } +} + +// TestDrainWorker marks a Worker as terminating, and checks the state reaches +// the worker cache the scheduler places actors from. Without that, a drained +// Worker would keep taking work. +func TestDrainWorker(t *testing.T) { + ns := namespaceForTest("ns-worker-drain") + tc := setupTest(t, ns) + defer tc.cleanup() + + registerWorker(t, tc, ns) + waitForWorkerAvailable(t, tc, testWorkerName) + + drained, err := tc.client.DrainWorker(context.Background(), &ateapipb.DrainWorkerRequest{Worker: workerRef(testWorkerName)}) + if err != nil { + t.Fatalf("DrainWorker failed: %v", err) + } + + want := newTestWorker(ns) + want.Metadata.Version = 2 + want.Status = &ateapipb.WorkerStatus{State: ateapipb.WorkerState_WORKER_STATE_DRAINING} + if diff := cmp.Diff(want, drained, protocmp.Transform(), ignoreServerMetadata); diff != "" { + t.Errorf("DrainWorker response mismatch (-want +got):\n%s", diff) + } + + waitForWorkerState(t, tc, testWorkerName, ateapipb.WorkerState_WORKER_STATE_DRAINING) +} + +// TestDeleteWorker deregisters a Worker and checks it is gone. Deregistering is +// not silently idempotent, so a second attempt reports NOT_FOUND. +func TestDeleteWorker(t *testing.T) { + ns := namespaceForTest("ns-worker-delete") + tc := setupTest(t, ns) + defer tc.cleanup() + ctx := context.Background() + + created := registerWorker(t, tc, ns) + + deleted, err := tc.client.DeleteWorker(ctx, &ateapipb.DeleteWorkerRequest{Worker: workerRef(testWorkerName)}) + if err != nil { + t.Fatalf("DeleteWorker failed: %v", err) + } + if diff := cmp.Diff(created, deleted, protocmp.Transform(), ignoreTimestamps); diff != "" { + t.Errorf("DeleteWorker returned something other than the record it removed (-created +deleted):\n%s", diff) + } + + listed, err := tc.client.ListWorkers(ctx, &ateapipb.ListWorkersRequest{}) + if err != nil { + t.Fatalf("ListWorkers failed: %v", err) + } + if containsWorker(listed.GetWorkers(), testWorkerName) { + t.Errorf("ListWorkers still returns worker %s after DeleteWorker", testWorkerName) + } + + _, err = tc.client.DeleteWorker(ctx, &ateapipb.DeleteWorkerRequest{Worker: workerRef(testWorkerName)}) + assertGrpcError(t, err, codes.NotFound, fmt.Sprintf("Worker %s not found", testWorkerName)) +} + func TestValidation_Worker(t *testing.T) { ns := namespaceForTest("ns-validation-worker") tc := setupTest(t, ns)