substrate/cmd/ateapi/internal/controlapi/functionaltest/worker_test.go

// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

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

// 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. Status is
// output-only, and that includes capacity — a Worker gets that from its own
// ateom's report to WorkerService, not from whoever registered it. These tests
// never place an actor, so the unreported default is enough.
//
// 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",
	}
}

// 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, registers
//     the Worker the syncer would derive from it, and reports the capacity its
//     ateom would.
//  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)
	defer tc.cleanup()

	createWorkerPool(t, tc, ns, "pool1", map[string]string{"foo": "bar"})
	podUID := createWorkerPod(t, tc, ns, "worker-1", "node1", "pool1")

	listResp, err := tc.client.ListWorkers(context.Background(), &ateapipb.ListWorkersRequest{})
	if err != nil {
		t.Fatalf("ListWorkers failed: %v", err)
	}

	var filteredWorkers []*ateapipb.Worker
	for _, w := range listResp.GetWorkers() {
		if w.GetWorkerNamespace() == ns {
			filteredWorkers = append(filteredWorkers, w)
		}
	}

	want := []*ateapipb.Worker{
		{
			Metadata: &ateapipb.ResourceMetadata{
				Name: podUID,
				// Two writes: the registration, then the capacity report.
				Version: 2,
			},
			WorkerNamespace: ns,
			WorkerPool:      "pool1",
			WorkerPod:       "worker-1",
			WorkerPodUid:    podUID,
			NodeName:        "node1",
			Ip:              "127.0.0.1",
			SandboxClass:    "gvisor",
			Labels:          map[string]string{"foo": "bar"},
			Status:          &ateapipb.WorkerStatus{State: ateapipb.WorkerState_WORKER_STATE_ACTIVE, Capacity: &ateapipb.WorkerResources{Actors: 1}},
		},
	}

	if diff := cmp.Diff(want, filteredWorkers, protocmp.Transform(), ignoreServerMetadata); diff != "" {
		t.Errorf("ListWorkers response mismatch (-want +got):\n%s", diff)
	}
}

// TestListWorkerActorAssignments reads the Actors a Worker hosts. They are a
// subresource rather than a field on Worker, so this is the only way to see
// them: a Worker hosting nothing lists an empty page, and one an Actor was
// placed on lists that Actor.
func TestListWorkerActorAssignments(t *testing.T) {
	ns := namespaceForTest("ns-worker-assignments")
	tc := setupTest(t, ns)
	defer tc.cleanup()
	ctx := context.Background()

	// createTemplate makes pool1 and tmpl1; the worker is what tmpl1's actors
	// get placed on.
	createTemplate(t, tc, ns)
	podUID := createWorkerPod(t, tc, ns, "worker-1", "node1", "pool1")

	empty, err := tc.client.ListWorkerActorAssignments(ctx, &ateapipb.ListWorkerActorAssignmentsRequest{Worker: workerRef(podUID)})
	if err != nil {
		t.Fatalf("ListWorkerActorAssignments on an empty worker failed: %v", err)
	}
	if got := len(empty.GetActorAssignments()); got != 0 {
		t.Errorf("worker hosting nothing lists %d assignments, want 0", got)
	}

	if _, err := tc.client.CreateActor(ctx, &ateapipb.CreateActorRequest{Actor: &ateapipb.Actor{
		Metadata:      &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: "id1"},
		ActorTemplate: &ateapipb.ObjectRef{Atespace: testAtespace, Name: "tmpl1"},
	}}); err != nil {
		t.Fatalf("CreateActor failed: %v", err)
	}
	if _, err := tc.client.ResumeActor(ctx, &ateapipb.ResumeActorRequest{
		Actor: &ateapipb.ObjectRef{Atespace: testAtespace, Name: "id1"},
	}); err != nil {
		t.Fatalf("ResumeActor failed: %v", err)
	}

	listed, err := tc.client.ListWorkerActorAssignments(ctx, &ateapipb.ListWorkerActorAssignmentsRequest{Worker: workerRef(podUID)})
	if err != nil {
		t.Fatalf("ListWorkerActorAssignments failed: %v", err)
	}
	if len(listed.GetActorAssignments()) != 1 {
		t.Fatalf("worker hosting one actor lists %d assignments, want 1", len(listed.GetActorAssignments()))
	}
	want := &ateapipb.ObjectRef{Atespace: testAtespace, Name: "id1"}
	if diff := cmp.Diff(want, listed.GetActorAssignments()[0].GetActor(), protocmp.Transform()); diff != "" {
		t.Errorf("assignment names the wrong actor (-want +got):\n%s", diff)
	}
}

// TestUpdateWorker changes labels, the only mutable field. An update replaces
// the whole Worker, so the request is the observed one with that field 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.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.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. The delete
// drains the Worker on its way out, so the record it returns is that one rather
// than the one Create stored. 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()

	registerWorker(t, tc, ns)

	deleted, err := tc.client.DeleteWorker(ctx, &ateapipb.DeleteWorkerRequest{Worker: workerRef(testWorkerName)})
	if err != nil {
		t.Fatalf("DeleteWorker 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, deleted, protocmp.Transform(), ignoreServerMetadata); diff != "" {
		t.Errorf("DeleteWorker response mismatch (-want +got):\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)
	defer tc.cleanup()

	t.Run("ListWorkers", func(t *testing.T) {
		_, err := tc.client.ListWorkers(context.Background(), &ateapipb.ListWorkersRequest{PageSize: -1})
		assertGrpcErrorRegex(t, err, codes.InvalidArgument, "page_size: Invalid value")
	})

	t.Run("ListWorkers invalid token", func(t *testing.T) {
		_, err := tc.client.ListWorkers(context.Background(), &ateapipb.ListWorkersRequest{PageToken: "%%%"})
		assertGrpcError(t, err, codes.InvalidArgument, "invalid page_token")
	})
}

func TestDeleteWorker_Preconditions(t *testing.T) {
	ns := namespaceForTest("ns-worker-delete-preconditions")
	tc := setupTest(t, ns)
	defer tc.cleanup()
	ctx := context.Background()
	ref := workerRef(testWorkerName)
	del := func(opts *ateapipb.DeleteOptions) error {
		_, err := tc.client.DeleteWorker(ctx, &ateapipb.DeleteWorkerRequest{Worker: ref, Options: opts})
		return err
	}

	worker := registerWorker(t, tc, ns)
	uid, version := worker.GetMetadata().GetUid(), worker.GetMetadata().GetVersion()

	assertGrpcError(t, del(&ateapipb.DeleteOptions{Version: version + 1}), codes.Aborted, "concurrent update conflict, please retry")
	assertGrpcError(t, del(&ateapipb.DeleteOptions{Uid: uid, Version: version + 1}), codes.Aborted, "concurrent update conflict, please retry")
	assertGrpcError(t, del(&ateapipb.DeleteOptions{Uid: foreignUID}), codes.Aborted, "Worker "+testWorkerName+" does not have uid "+foreignUID)
	assertGrpcError(t, del(&ateapipb.DeleteOptions{Uid: foreignUID, Version: version}), codes.Aborted, "Worker "+testWorkerName+" does not have uid "+foreignUID)
	if _, err := tc.client.GetWorker(ctx, &ateapipb.GetWorkerRequest{Worker: ref}); err != nil {
		t.Fatalf("a refused delete removed the worker: %v", err)
	}

	// The delete drains the worker first, which moves the version; the guard
	// is checked against the record the caller read.
	if err := del(&ateapipb.DeleteOptions{Version: version}); err != nil {
		t.Fatalf("DeleteWorker with the matching version: %v", err)
	}
	worker = registerWorker(t, tc, ns)
	if err := del(&ateapipb.DeleteOptions{Uid: worker.GetMetadata().GetUid()}); err != nil {
		t.Fatalf("DeleteWorker with the matching uid: %v", err)
	}
	worker = registerWorker(t, tc, ns)
	if err := del(&ateapipb.DeleteOptions{Uid: worker.GetMetadata().GetUid(), Version: worker.GetMetadata().GetVersion()}); err != nil {
		t.Fatalf("DeleteWorker with both guards: %v", err)
	}
	_, err := tc.client.GetWorker(ctx, &ateapipb.GetWorkerRequest{Worker: ref})
	assertGrpcError(t, err, codes.NotFound, "Worker "+testWorkerName+" not found")
}