mirror of
https://github.com/agent-substrate/substrate.git
synced 2026-11-10 10:13:50 +08:00
sfunkenhauser
8a6b5e149a
Fixes #944 - [x] Tests pass - [x] Appropriate changes to documentation are included in the PR
365 lines
14 KiB
Go
365 lines
14 KiB
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",
|
|
Ips: []string{"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",
|
|
Ips: []string{"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")
|
|
}
|