substrate/cmd/ateapi/internal/controlapi/functionaltest/common_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"
	"net"
	"regexp"
	"testing"
	"time"

	"go.ate.dev/substrate/cmd/ateapi/internal/controlapi"
	"go.ate.dev/substrate/cmd/ateapi/internal/store"
	"go.ate.dev/substrate/cmd/ateapi/internal/store/storetest"
	"go.ate.dev/substrate/cmd/ateapi/internal/workercache"
	"go.ate.dev/substrate/internal/ateinterceptors"
	"go.ate.dev/substrate/internal/installdefaults"
	"go.ate.dev/substrate/internal/localca"
	"go.ate.dev/substrate/internal/localjwtauthority"
	"go.ate.dev/substrate/internal/objectstore/objectstoretest"
	"go.ate.dev/substrate/internal/objectstoreplugin/objectstoreplugintest"
	"go.ate.dev/substrate/internal/preview"
	"go.ate.dev/substrate/internal/resources"
	"go.ate.dev/substrate/internal/volume"
	atev1alpha1 "go.ate.dev/substrate/pkg/api/v1alpha1"
	"go.ate.dev/substrate/pkg/client/clientset/versioned"
	"go.ate.dev/substrate/pkg/client/informers/externalversions"
	"go.ate.dev/substrate/pkg/proto/ateapipb"
	"go.opentelemetry.io/otel/attribute"
	sdkmetric "go.opentelemetry.io/otel/sdk/metric"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/sdk/trace/tracetest"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/status"
	"google.golang.org/grpc/test/bufconn"
	"google.golang.org/protobuf/testing/protocmp"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/cache"
	"k8s.io/utils/ptr"
)

const (
	testAtespace = "test-atespace"
	testActorID  = "id1"
	// testStorageLocation is the snapshot_config.storage_location the test
	// templates hand out. No object store is wired up behind it.
	testStorageLocation = "gs://fake-fake-fake"
	testActorJWTIssuer  = "https://idp.example.test"

	// ateletNamespace and byNode mirror the unexported constants controlapi's
	// atelet informer is built with.
	ateletNamespace = "ate-system"
	byNode          = "by-node"
)

var (
	ignoreUID        = protocmp.IgnoreFields(&ateapipb.ResourceMetadata{}, "uid")
	ignoreVersion    = protocmp.IgnoreFields(&ateapipb.ResourceMetadata{}, "version")
	ignoreTimestamps = protocmp.IgnoreFields(&ateapipb.ResourceMetadata{}, "create_time", "update_time")

	// ignoreServerMetadata skips the ResourceMetadata fields the store assigns.
	ignoreServerMetadata = protocmp.IgnoreFields(&ateapipb.ResourceMetadata{}, "uid", "create_time", "update_time")
)

type testContext struct {
	service         *controlapi.RPCService
	client          ateapipb.ControlClient
	noOptInClient   ateapipb.ControlClient // never sends preview.Header
	k8sClient       kubernetes.Interface
	substrateClient versioned.Interface
	persistence     store.Interface
	workerCache     *workercache.Cache
	fakeAtelet      *FakeAteletServer
	cleanup         func()
	// ateletIndexer is the index DialForAteletOnNode looks up atelets in.
	// setupAteletOnNode waits on it so a test never dials a node whose atelet the
	// informer has not seen yet.
	ateletIndexer cache.Indexer
	// metricReader collects what the service's instruments recorded.
	metricReader *sdkmetric.ManualReader
	// objectStore holds the external snapshots the service wrote, so a test can
	// see which ones a flow created and released.
	objectStore *objectstoretest.Fake
}

// setupTest sets up a fully isolated test environment.
func setupTest(t *testing.T, ns string) *testContext {
	t.Helper()
	return setupTestWithOptions(t, ns, testOptions{})
}

type testOptions struct {
	allowPreview  bool // enable use of preview features over gRPC
	volumePlugins map[string]volume.VolumePluginControlPlane
}

// setupPreviewTest is setupTest with a client that opts in to preview
// features. It does not enable any gate; call preview.InitForTest for that.
func setupPreviewTest(t *testing.T, ns string) *testContext {
	t.Helper()
	return setupTestWithOptions(t, ns, testOptions{allowPreview: true})
}

// setupTestWithVolumePlugins is setupPreviewTest with the default mock volume
// plugin replaced by plugins, keyed by driver name. Tests that need a
// failure-injecting plugin pass it here rather than swapping it into the
// running RPCService, so each test owns its own plugin set.
func setupTestWithVolumePlugins(t *testing.T, ns string, plugins map[string]volume.VolumePluginControlPlane) *testContext {
	t.Helper()
	return setupTestWithOptions(t, ns, testOptions{volumePlugins: plugins, allowPreview: true})
}

func setupTestWithOptions(t *testing.T, ns string, options testOptions) *testContext {
	t.Helper()

	var dialOpts []grpc.DialOption
	if options.allowPreview {
		dialOpts = preview.DialOptionsWithOptIn()
	}

	// 1. Start an isolated PostgreSQL-backed store.
	persistence, cleanupStore := storetest.SetupTestStore(t)

	// 2. Initialize Clientsets using global cfg
	k8sClient, err := kubernetes.NewForConfig(cfg)
	if err != nil {
		cleanupStore()
		t.Fatalf("failed to create k8s clientset: %v", err)
	}

	substrateClient, err := versioned.NewForConfig(cfg)
	if err != nil {
		cleanupStore()
		t.Fatalf("failed to create substrate clientset: %v", err)
	}

	// 3. Initialize Informers
	ateletFactory, ateletInformer := controlapi.AteletInformer(k8sClient, installdefaults.SystemNamespace)
	scFactory := informers.NewSharedInformerFactory(k8sClient, 0)
	scLister := scFactory.Storage().V1().StorageClasses().Lister()

	substrateInformerFactory := externalversions.NewSharedInformerFactory(substrateClient, 0)
	csiDriverConfigLister := substrateInformerFactory.Api().V1alpha1().CSIDriverConfigs().Lister()

	ctx, cancel := context.WithCancel(context.Background())

	ateletFactory.Start(ctx.Done())
	substrateInformerFactory.Start(ctx.Done())
	scFactory.Start(ctx.Done())

	ateletFactory.WaitForCacheSync(ctx.Done())
	substrateInformerFactory.WaitForCacheSync(ctx.Done())
	scFactory.WaitForCacheSync(ctx.Done())

	// 4. Initialize Service
	wc := workercache.New(persistence, 5*time.Minute)
	if err := wc.Start(ctx); err != nil {
		cancel()
		cleanupStore()
		t.Fatalf("failed to start worker cache: %v", err)
	}

	// Dial the fake atelet over insecure transport instead of per-atelet mTLS,
	// so DialForAteletOnNode's real lookup/dial/cache path is exercised under test.
	dialer := controlapi.NewAteletDialer(ateletInformer.GetIndexer(), installdefaults.AteletSPIFFEID(installdefaults.SystemNamespace), "", "",
		controlapi.WithDialCredentials(func(_ string) (credentials.TransportCredentials, error) {
			return insecure.NewCredentials(), nil
		}))

	metricReader := sdkmetric.NewManualReader()
	instruments, err := controlapi.NewInstruments(sdkmetric.NewMeterProvider(sdkmetric.WithReader(metricReader)).Meter("ateapi"))
	if err != nil {
		cancel()
		cleanupStore()
		t.Fatalf("failed to create metric instruments: %v", err)
	}
	volPlugins := options.volumePlugins
	if volPlugins == nil {
		mockPlugin := volume.NewMockVolumePlugin()
		mockDriverName, err := mockPlugin.DriverName(ctx)
		if err != nil {
			t.Fatalf("failed to get mock driver name: %v", err)
		}
		volPlugins = map[string]volume.VolumePluginControlPlane{
			mockDriverName: mockPlugin,
		}
	}

	actorJWTAuthority, err := localjwtauthority.GenerateAuthority("ES256", "1")
	if err != nil {
		t.Fatalf("Error generating actor JWT authority: %v", err)
	}

	actorJWTAuthorityPool := &localjwtauthority.ConcretePool{
		Authorities: []*localjwtauthority.Authority{
			actorJWTAuthority,
		},
		ActiveForSigning: "1",
	}

	actorCA, err := localca.GenerateCA("1", localca.KeyTypeECDSAP256, 365*24*time.Hour)
	if err != nil {
		t.Fatalf("Error generating actor CA: %v", err)
	}

	actorCAPool := &localca.ConcretePool{
		CAs: []*localca.CA{
			actorCA,
		},
		ActiveForSigning: "1",
	}

	objectStore := objectstoretest.New()

	service := controlapi.NewRPCService(
		persistence,
		wc,
		csiDriverConfigLister,
		scLister,
		dialer,
		instruments,
		"",
		volPlugins,
		objectstoreplugintest.ControlClient(objectStore),
		testActorJWTIssuer,
		actorJWTAuthorityPool,
		actorCAPool,
		testStorageLocation,
	)

	// 5. Start REAL gRPC Server for ATE API
	grpcServer := grpc.NewServer(grpc.ChainUnaryInterceptor(
		ateinterceptors.ServerUnaryInterceptor,
		ateinterceptors.RejectUnknownFieldsUnaryInterceptor,
		preview.UnaryServerInterceptor,
	))
	ateapipb.RegisterControlServer(grpcServer, service)

	lis := bufconn.Listen(8192)

	go func() {
		if err := grpcServer.Serve(lis); err != nil {
			t.Logf("grpc server exited: %v", err)
		}
	}()

	baseDialOpts := []grpc.DialOption{
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithContextDialer(func(ctx context.Context, s string) (net.Conn, error) {
			return lis.Dial()
		}),
	}
	conn, err := grpc.NewClient("passthrough://bufnet", append(baseDialOpts, dialOpts...)...)
	if err != nil {
		grpcServer.Stop()
		cancel()
		cleanupStore()
		t.Fatalf("failed to connect: %v", err)
	}
	noOptInConn, err := grpc.NewClient("passthrough://bufnet", baseDialOpts...)
	if err != nil {
		conn.Close()
		grpcServer.Stop()
		cancel()
		cleanupStore()
		t.Fatalf("failed to connect: %v", err)
	}
	t.Cleanup(func() { noOptInConn.Close() })

	client := ateapipb.NewControlClient(conn)

	// Call Reset on global mock
	fakeAtelet.Reset()
	// A checkpoint writes its external snapshot into this test's own store, so
	// the copy and release steps have something to act on.
	fakeAtelet.SetObjectStore(objectStore)

	// Create namespace
	_, err = k8sClient.CoreV1().Namespaces().Create(context.Background(), &corev1.Namespace{
		ObjectMeta: metav1.ObjectMeta{Name: ns},
	}, metav1.CreateOptions{})
	if err != nil {
		conn.Close()
		grpcServer.Stop()
		cancel()
		cleanupStore()
		t.Fatalf("failed to create namespace %s: %v", ns, err)
	}

	// CreateActor now requires the atespace to exist first.
	if _, err := client.CreateAtespace(context.Background(), &ateapipb.CreateAtespaceRequest{Atespace: &ateapipb.Atespace{Metadata: &ateapipb.ResourceMetadata{Name: testAtespace}}}); err != nil {
		conn.Close()
		grpcServer.Stop()
		cancel()
		cleanupStore()
		t.Fatalf("failed to seed test atespace %q: %v", testAtespace, err)
	}

	cleanup := func() {
		conn.Close()
		grpcServer.Stop()
		cancel()
		cleanupStore()
	}

	return &testContext{
		service:         service,
		client:          client,
		noOptInClient:   ateapipb.NewControlClient(noOptInConn),
		k8sClient:       k8sClient,
		substrateClient: substrateClient,
		persistence:     persistence,
		workerCache:     wc,
		fakeAtelet:      fakeAtelet,
		cleanup:         cleanup,
		ateletIndexer:   ateletInformer.GetIndexer(),
		metricReader:    metricReader,
		objectStore:     objectStore,
	}
}

// snapshotObjectNames returns the names, relative to snapshotURI, of the
// objects the external snapshot there is made of. Empty means the snapshot is
// not in object storage — either never written, or collected.
func snapshotObjectNames(t *testing.T, tc *testContext, snapshotURI string) []string {
	t.Helper()
	uri, err := resources.ParseSnapshotURI(snapshotURI)
	if err != nil {
		t.Fatalf("ParseSnapshotURI(%q) = %v", snapshotURI, err)
	}
	return tc.objectStore.Snapshot(t, uri)
}

// assertSnapshotPresent fails when the external snapshot at snapshotURI is not
// in object storage.
func assertSnapshotPresent(t *testing.T, tc *testContext, snapshotURI string) {
	t.Helper()
	if names := snapshotObjectNames(t, tc, snapshotURI); len(names) == 0 {
		t.Errorf("external snapshot %s is not in object storage, want it present", snapshotURI)
	}
}

// assertSnapshotCollected fails when anything is left of the external snapshot
// at snapshotURI.
func assertSnapshotCollected(t *testing.T, tc *testContext, snapshotURI string) {
	t.Helper()
	if names := snapshotObjectNames(t, tc, snapshotURI); len(names) != 0 {
		t.Errorf("external snapshot %s still holds %v, want it collected", snapshotURI, names)
	}
}

const goldenSnapshotName = "9c2f7b41-6d05-4e83-a1f7-3b8c0d5e2a94"

// goldenSnapshotURI is the snapshot owned by the test template's golden tag.
func goldenSnapshotURI(t *testing.T) string {
	t.Helper()
	uri, err := resources.NewTagSnapshotURI(testStorageLocation, testAtespace, goldenSnapshotName)
	if err != nil {
		t.Fatalf("NewTagSnapshotURI: %v", err)
	}
	return uri.String()
}

// snapshotOwnedByActor reports whether snapshotURI sits under the actor's own
// prefix. That is what says the actor took the snapshot itself rather than
// borrowing it from a tag, and so what decides whether its suspends and its
// delete may collect it.
func snapshotOwnedByActor(t *testing.T, actor *ateapipb.Actor, snapshotURI string) bool {
	t.Helper()
	uri, err := resources.ParseSnapshotURI(snapshotURI)
	if err != nil {
		t.Fatalf("ParseSnapshotURI(%q) = %v", snapshotURI, err)
	}
	return uri.OwnedBy(resources.ActorSnapshotOwner(actor.GetMetadata().GetAtespace(), actor.GetMetadata().GetUid()))
}

// assertSnapshotOwnedByActor fails when snapshotURI is not under the actor's
// own prefix.
func assertSnapshotOwnedByActor(t *testing.T, actor *ateapipb.Actor, snapshotURI string) {
	t.Helper()
	if !snapshotOwnedByActor(t, actor, snapshotURI) {
		t.Errorf("external snapshot %s is not under actor %s's own prefix, want it owned by the actor", snapshotURI, resources.ActorRefFromActor(actor))
	}
}

// snapshotName returns the name snapshotURI gives its snapshot in object
// storage, which the Snapshot's uuid must match.
func snapshotName(t *testing.T, snapshotURI string) string {
	t.Helper()
	uri, err := resources.ParseSnapshotURI(snapshotURI)
	if err != nil {
		t.Fatalf("ParseSnapshotURI(%q) = %v", snapshotURI, err)
	}
	return uri.Name()
}

func namespaceForTest(baseName string) string {
	return fmt.Sprintf("%s-%d", baseName, time.Now().UnixNano())
}

func createTemplate(t *testing.T, tc *testContext, ns string) *ateapipb.ActorTemplate {
	t.Helper()
	return createTemplateWithContainers(t, tc, ns, []*ateapipb.Container{
		{
			Name:    "main",
			Image:   "main@sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
			Command: []string{"/main"},
		},
	})
}

// createAtespace creates an atespace via the API.
func createAtespace(t *testing.T, tc *testContext, name string) {
	t.Helper()
	if _, err := tc.client.CreateAtespace(context.Background(), &ateapipb.CreateAtespaceRequest{Atespace: &ateapipb.Atespace{Metadata: &ateapipb.ResourceMetadata{Name: name}}}); err != nil {
		t.Fatalf("CreateAtespace(%s) failed: %v", name, err)
	}
}

const poolLabelKey = "pool"

func createTemplateWithContainers(t *testing.T, tc *testContext, ns string, containers []*ateapipb.Container) *ateapipb.ActorTemplate {
	t.Helper()
	return createTemplateWithContainersAndVolumes(t, tc, ns, containers, nil)
}

func createTemplateWithVolumes(t *testing.T, tc *testContext, ns string, volumes []*ateapipb.Volume, mounts []*ateapipb.VolumeMount) *ateapipb.ActorTemplate {
	t.Helper()
	return createTemplateWithContainersAndVolumes(t, tc, ns, []*ateapipb.Container{
		{
			Name:         "main",
			Image:        "main@sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
			Command:      []string{"/main"},
			VolumeMounts: mounts,
		},
	}, volumes)
}

// createTemplateWithContainersAndVolumes creates the substrate ActorTemplate
// "tmpl1" in testAtespace, backed by a WorkerPool in ns whose labels match the
// template's worker selector, and seeds its golden snapshot. ns keys the pool
// labels, so each test's template still selects only its own pool.
func createTemplateWithContainersAndVolumes(t *testing.T, tc *testContext, ns string, containers []*ateapipb.Container, volumes []*ateapipb.Volume) *ateapipb.ActorTemplate {
	t.Helper()

	createWorkerPool(t, tc, ns, "pool1", map[string]string{poolLabelKey: ns})

	created, err := tc.client.CreateActorTemplate(context.Background(), &ateapipb.CreateActorTemplateRequest{
		ActorTemplate: &ateapipb.ActorTemplate{
			Metadata: &ateapipb.ResourceMetadata{
				Atespace: testAtespace,
				Name:     "tmpl1",
			},
			SnapshotConfig: &ateapipb.SnapshotConfig{
				StorageLocation: testStorageLocation,
			},
			Containers: containers,
			Volumes:    volumes,
			WorkerSelector: &ateapipb.Selector{
				MatchLabels: map[string]string{poolLabelKey: ns},
			},
			Resources: &ateapipb.ResourceRequirements{Limits: resources.CPUMemory(1000, 512<<20)},
		},
	})
	if err != nil {
		t.Fatalf("failed to create actor template: %v", err)
	}

	tag, err := tc.persistence.CreateTag(context.Background(), &ateapipb.Tag{
		Metadata:    &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: created.GetMetadata().GetUid()},
		SourceActor: &ateapipb.ObjectRef{Atespace: testAtespace, Name: created.GetMetadata().GetUid()},
		Scope:       ateapipb.TagScope_TAG_SCOPE_PUBLISHED,
		Status: &ateapipb.TagStatus{
			Snapshot:         newDurableSnapshot(1, ateapipb.SnapshotOwner_SNAPSHOT_OWNER_TAG, ateapipb.SnapshotFidelity_SNAPSHOT_FIDELITY_MEMORY, created.GetMetadata().GetUid(), goldenSnapshotName, goldenSnapshotURI(t)),
			ActorTemplateUid: created.GetMetadata().GetUid(),
		},
	})
	if err != nil {
		t.Fatalf("create golden tag: %v", err)
	}

	// Record the golden snapshot on the template's status directly in the
	// store, as the ActorTemplateReconciler's checkpoint would: there is no
	// status RPC, and the reconciler does not run in this test environment.
	updated, err := tc.persistence.UpdateActorTemplate(context.Background(),
		resources.ActorTemplateRefFromActorTemplate(created), store.PreconditionFrom(created),
		func(dbTemplate *ateapipb.ActorTemplate) error {
			dbTemplate.Status = &ateapipb.ActorTemplateStatus{
				GoldenSnapshotStatus: &ateapipb.GoldenSnapshotStatus{
					GoldenTag: resources.TagRefFromTag(tag).ToObjectRef(),
				},
			}
			return nil
		})
	if err != nil {
		t.Fatalf("failed to record the template's golden snapshot: %v", err)
	}
	return updated
}

func newDurableSnapshot(gen int32, owner ateapipb.SnapshotOwner, fidelity ateapipb.SnapshotFidelity, templateUID, uuid, uri string) *ateapipb.Snapshot {
	return &ateapipb.Snapshot{
		Uuid:             uuid,
		Generation:       gen,
		Owner:            owner,
		Fidelity:         fidelity,
		ActorTemplateUid: templateUID,
		DurableSnapshot: &ateapipb.SnapshotStorage{
			Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED,
			Object: &ateapipb.ObjectSnapshot{SnapshotUri: uri},
		},
	}
}

// newPausedSnapshot is the snapshot a pause leaves on node locality: its
// durable copy, under the actor's own prefix, is PENDING until atelet uploads it.
func newPausedSnapshot(t *testing.T, gen int32, fidelity ateapipb.SnapshotFidelity, templateUID, actorUID, uuid, locality string) *ateapipb.Snapshot {
	t.Helper()
	uri, err := resources.NewActorSnapshotURI(testStorageLocation, testAtespace, actorUID, uuid)
	if err != nil {
		t.Fatalf("NewActorSnapshotURI: %v", err)
	}
	return &ateapipb.Snapshot{
		Uuid:             uuid,
		Generation:       gen,
		Owner:            ateapipb.SnapshotOwner_SNAPSHOT_OWNER_ACTOR,
		Fidelity:         fidelity,
		ActorTemplateUid: templateUID,
		Locality:         locality,
		DurableSnapshot: &ateapipb.SnapshotStorage{
			Status: ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_PENDING,
			Object: &ateapipb.ObjectSnapshot{SnapshotUri: uri.String()},
		},
	}
}

func durableSnapshot(status *ateapipb.ActorStatus) *ateapipb.Snapshot {
	var best *ateapipb.Snapshot
	for _, snap := range status.GetSnapshots() {
		st := snap.GetDurableSnapshot()
		if st.GetStatus() == ateapipb.SnapshotStorageStatus_SNAPSHOT_STORAGE_STATUS_COMPLETED &&
			st.GetObject().GetSnapshotUri() != "" {
			if best == nil || snap.GetGeneration() > best.GetGeneration() {
				best = snap
			}
		}
	}
	return best
}

func durableSnapshotStorage(status *ateapipb.ActorStatus) *ateapipb.SnapshotStorage {
	return durableSnapshot(status).GetDurableSnapshot()
}

func durableSnapshotURI(status *ateapipb.ActorStatus) string {
	return durableSnapshotStorage(status).GetObject().GetSnapshotUri()
}

func localSnapshot(status *ateapipb.ActorStatus) *ateapipb.Snapshot {
	var best *ateapipb.Snapshot
	for _, snap := range status.GetSnapshots() {
		if snap.GetLocality() != "" && snap.GetUuid() != "" {
			if best == nil || snap.GetGeneration() > best.GetGeneration() {
				best = snap
			}
		}
	}
	return best
}

func tagSnapshotURI(tag *ateapipb.Tag) string {
	return tag.GetStatus().GetSnapshot().GetDurableSnapshot().GetObject().GetSnapshotUri()
}

func createWorkerPool(t *testing.T, tc *testContext, ns string, name string, labels map[string]string) {
	t.Helper()
	wp := &atev1alpha1.WorkerPool{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: ns,
			Labels:    labels,
		},
		Spec: atev1alpha1.WorkerPoolSpec{
			Replicas:       1,
			WorkerImage:    "worker@sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
			SandboxClasses: []atev1alpha1.WorkerPoolSandboxClass{{Name: atev1alpha1.SandboxClassGvisor}},
		},
	}
	if _, err := tc.substrateClient.ApiV1alpha1().WorkerPools(ns).Create(context.Background(), wp, metav1.CreateOptions{}); err != nil {
		t.Fatalf("failed to create WorkerPool: %v", err)
	}
}

// createTemplateWithSelector creates a substrate ActorTemplate in
// testAtespace with the given worker selector and no golden snapshot.
func createTemplateWithSelector(t *testing.T, tc *testContext, name string, selector *ateapipb.Selector) *ateapipb.ActorTemplate {
	t.Helper()
	created, err := tc.client.CreateActorTemplate(context.Background(), &ateapipb.CreateActorTemplateRequest{
		ActorTemplate: &ateapipb.ActorTemplate{
			Metadata: &ateapipb.ResourceMetadata{
				Atespace: testAtespace,
				Name:     name,
			},
			SnapshotConfig: &ateapipb.SnapshotConfig{
				StorageLocation: testStorageLocation,
			},
			Containers: []*ateapipb.Container{
				{Name: "main", Image: "main@sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", Command: []string{"/main"}},
			},
			WorkerSelector: selector,
			Resources:      &ateapipb.ResourceRequirements{Limits: resources.CPUMemory(1000, 512<<20)},
		},
	})
	if err != nil {
		t.Fatalf("failed to create actor template: %v", err)
	}
	return created
}

// createWorkerPod creates a worker pod, registers the matching Worker, and
// waits for it to reach the worker cache. It returns the name of the resulting
// Worker, which is the key to look it up by.
//
// Both halves are needed: the Worker record is what the scheduler places actors
// on, and the pod is what the atelet dialer resolves to reach one. In a real
// cluster the syncer in ate-controller derives the first from the second; here
// the test writes both, building the Worker exactly as the syncer would.
func createWorkerPod(t *testing.T, tc *testContext, ns string, name string, nodeName string, poolName string) string {
	t.Helper()
	pod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: ns,
			Labels: map[string]string{
				"ate.dev/worker-pool": poolName,
			},
		},
		Spec: corev1.PodSpec{
			NodeName: nodeName,
			Containers: []corev1.Container{
				{Name: "main", Image: "nginx"},
			},
		},
	}
	createdPod, err := tc.k8sClient.CoreV1().Pods(ns).Create(context.Background(), pod, metav1.CreateOptions{})
	if err != nil {
		t.Fatalf("failed to create worker pod: %v", err)
	}
	createdPod.Status.PodIPs = []corev1.PodIP{{IP: "127.0.0.1"}}
	createdPod.Status.Phase = corev1.PodRunning
	createdPod.Status.Conditions = []corev1.PodCondition{{
		Type:   corev1.PodReady,
		Status: corev1.ConditionTrue,
	}}
	_, err = tc.k8sClient.CoreV1().Pods(ns).UpdateStatus(context.Background(), createdPod, metav1.UpdateOptions{})
	if err != nil {
		t.Fatalf("failed to update worker pod status: %v", err)
	}

	// The WorkerPool supplies the fields the syncer copies off it.
	pool, err := tc.substrateClient.ApiV1alpha1().WorkerPools(ns).Get(context.Background(), poolName, metav1.GetOptions{})
	if err != nil {
		t.Fatalf("failed to get WorkerPool %s/%s: %v", ns, poolName, err)
	}
	if _, err := tc.client.CreateWorker(context.Background(), &ateapipb.CreateWorkerRequest{
		Worker: &ateapipb.Worker{
			// Workers are global-scoped and named after the pod UID.
			Metadata:        &ateapipb.ResourceMetadata{Name: string(createdPod.UID)},
			WorkerNamespace: ns,
			WorkerPool:      poolName,
			WorkerPod:       name,
			WorkerPodUid:    string(createdPod.UID),
			Ips:             []string{"127.0.0.1"},
			NodeName:        nodeName,
			SandboxClass:    string(pool.Spec.DefaultSandboxClass()),
			Labels:          pool.GetLabels(),
			// Capacity is not settable here: a Worker gets it from its own
			// worker's report, which the reportWorkerCapacity below stands in
			// for. These pods declare no limits, so only the actor ceiling is
			// reported; see setWorkerActorCapacity for tests needing more.
		},
	}); err != nil {
		t.Fatalf("failed to register worker: %v", err)
	}
	// As the syncer does once the pod is Ready.
	if _, err := tc.client.MarkWorkerAvailable(context.Background(), &ateapipb.MarkWorkerAvailableRequest{
		Worker: &ateapipb.ObjectRef{Name: string(createdPod.UID)},
	}); err != nil {
		t.Fatalf("failed to mark worker available: %v", err)
	}
	reportWorkerCapacity(t, tc, string(createdPod.UID), 1)

	// Wait for the worker to appear in worker cache.
	err = wait.PollUntilContextTimeout(context.Background(), 10*time.Millisecond, 5*time.Second, true, func(ctx context.Context) (bool, error) {
		workers, err := tc.workerCache.Workers()
		if err != nil {
			return false, nil // Cache not ready yet; retry.
		}
		for _, w := range workers {
			if w.GetWorkerNamespace() == ns && w.GetWorkerPod() == name {
				return true, nil
			}
		}
		return false, nil
	})
	if err != nil {
		t.Fatalf("failed to wait for worker to appear in worker cache: %v", err)
	}
	return string(createdPod.UID)
}

// waitForWorkerAvailable waits for an assignment release to reach the worker
// cache. Lifecycle RPCs commit the release to the store before the cache's
// PostgreSQL watch processes the corresponding update.
func waitForWorkerAvailable(t *testing.T, tc *testContext, workerName string) {
	t.Helper()
	err := wait.PollUntilContextTimeout(context.Background(), 10*time.Millisecond, 5*time.Second, true, func(context.Context) (bool, error) {
		worker, err := tc.workerCache.Worker(workerName)
		if err != nil {
			return false, nil
		}
		// Hosting nothing is what "available" means, and the allocation total
		// is how a cached Worker reports it: it does not carry the records.
		return worker.GetStatus().GetState() == ateapipb.WorkerState_WORKER_STATE_ACTIVE &&
			worker.GetStatus().GetAllocated().GetActors() == 0, nil
	})
	if err != nil {
		t.Fatalf("failed to wait for worker %s to become available: %v", workerName, err)
	}
}

// createAteletPod creates an atelet pod on nodeName and marks it Running with
// an IP, which DialForAteletOnNode requires. The pod carries the namespace and
// app=atelet label AteletInformer selects on, and is indexed by its node.
func createAteletPod(kc kubernetes.Interface, name, nodeName string) error {
	pod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: ateletNamespace,
			Labels:    map[string]string{"app": "atelet"},
		},
		Spec: corev1.PodSpec{
			NodeName:   nodeName,
			Containers: []corev1.Container{{Name: "main", Image: "nginx"}},
		},
	}
	created, err := kc.CoreV1().Pods(ateletNamespace).Create(context.Background(), pod, metav1.CreateOptions{})
	if apierrors.IsAlreadyExists(err) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("creating atelet pod %s on %s: %w", name, nodeName, err)
	}
	created.Status.PodIPs = []corev1.PodIP{{IP: "127.0.0.1"}}
	created.Status.Phase = corev1.PodRunning
	if _, err := kc.CoreV1().Pods(ateletNamespace).UpdateStatus(context.Background(), created, metav1.UpdateOptions{}); err != nil {
		return fmt.Errorf("updating atelet pod %s status: %w", name, err)
	}
	return nil
}

// setupAteletOnNode makes nodeName dialable for the duration of the test: it
// creates an atelet pod there, waits for the dialer's index to see it, and
// removes it on cleanup. The package fixture only creates one atelet, on
// node1, and the dialer resolves atelets per node, so a worker on any other
// node is unreachable without this.
func setupAteletOnNode(t *testing.T, tc *testContext, name, nodeName string) {
	t.Helper()
	if err := createAteletPod(tc.k8sClient, name, nodeName); err != nil {
		t.Fatalf("%v", err)
	}
	t.Cleanup(func() {
		_ = tc.k8sClient.CoreV1().Pods(ateletNamespace).Delete(context.Background(), name, metav1.DeleteOptions{
			GracePeriodSeconds: ptr.To[int64](0),
		})
	})

	// Wait for the dialer's index to see the newly created atlet
	err := wait.PollUntilContextTimeout(context.Background(), 10*time.Millisecond, 5*time.Second, true, func(ctx context.Context) (bool, error) {
		atelets, err := tc.ateletIndexer.ByIndex(byNode, nodeName)
		if err != nil {
			return false, nil
		}
		for _, obj := range atelets {
			p := obj.(*corev1.Pod)
			if p.Name == name && len(p.Status.PodIPs) > 0 {
				return true, nil
			}
		}
		return false, nil
	})
	if err != nil {
		t.Fatalf("failed to wait for atelet pod on %s to be indexed: %v", nodeName, err)
	}
}

func deleteWorkerPod(t *testing.T, tc *testContext, ns string, name string) {
	t.Helper()
	err := tc.k8sClient.CoreV1().Pods(ns).Delete(context.Background(), name, metav1.DeleteOptions{
		GracePeriodSeconds: ptr.To[int64](0),
	})
	if err != nil {
		t.Fatalf("failed to delete worker pod %s: %v", name, err)
	}

	// Deregister the Worker, as the syncer would once it saw the pod go. The
	// Worker is named after the pod UID, but the caller only has the pod name,
	// so it is found by the pod fields it carries.
	resp, err := tc.client.ListWorkers(context.Background(), &ateapipb.ListWorkersRequest{})
	if err != nil {
		t.Fatalf("failed to list workers: %v", err)
	}
	for _, w := range resp.GetWorkers() {
		if w.GetWorkerNamespace() != ns || w.GetWorkerPod() != name {
			continue
		}
		if _, err := tc.client.DeleteWorker(context.Background(), &ateapipb.DeleteWorkerRequest{
			Worker: &ateapipb.ObjectRef{Name: w.GetMetadata().GetName()},
		}); err != nil {
			t.Fatalf("failed to deregister worker %s: %v", w.GetMetadata().GetName(), err)
		}
	}

	err = wait.PollUntilContextTimeout(context.Background(), 10*time.Millisecond, 5*time.Second, true, func(ctx context.Context) (bool, error) {
		workers, err := tc.workerCache.Workers()
		if err != nil {
			return false, nil // Cache not ready yet; retry.
		}
		for _, w := range workers {
			if w.GetWorkerNamespace() == ns && w.GetWorkerPod() == name {
				return false, nil // Still there
			}
		}
		return true, nil
	})
	if err != nil {
		t.Fatalf("failed to wait for worker to be removed from worker cache: %v", err)
	}
}

func assertGrpcErrorRegex(t *testing.T, err error, wantCode codes.Code, wantMsg string) {
	t.Helper()
	fn := func(got string) (string, bool) {
		matched, matchErr := regexp.MatchString(wantMsg, got)
		if matchErr != nil {
			t.Fatalf("failed to compile regex %q: %v", wantMsg, matchErr)
		}
		return wantMsg, matched
	}
	assertGrpcErrorImpl(t, err, wantCode, fn)
}

func assertGrpcError(t *testing.T, err error, wantCode codes.Code, wantMsg string) {
	t.Helper()
	fn := func(got string) (string, bool) {
		return wantMsg, got == wantMsg
	}
	assertGrpcErrorImpl(t, err, wantCode, fn)
}

func assertGrpcErrorImpl(t *testing.T, err error, wantCode codes.Code, msgMatches func(got string) (string, bool)) {
	t.Helper()
	if err == nil {
		t.Fatalf("expected error, got nil")
	}
	st, ok := status.FromError(err)
	if !ok {
		t.Fatalf("expected gRPC status error, got: %v", err)
	}
	if st.Code() != wantCode {
		t.Errorf("expected status %v, got %v", wantCode, st.Code())
	}
	if want, ok := msgMatches(st.Message()); !ok {
		t.Errorf("expected message %q, got %q", want, st.Message())
	}
}

// recordRootSpanAttrs runs fn under a fresh recording root span from a local
// TracerProvider and returns that span's attributes, so a test can observe what
// the code under test stamps on the span carried in ctx. It never swaps the
// global provider (the code under test reads its span via trace.SpanFromContext,
// not the global provider), so span tests stay parallel-safe.
func recordRootSpanAttrs(t *testing.T, fn func(ctx context.Context)) map[attribute.Key]attribute.Value {
	t.Helper()
	sr := tracetest.NewSpanRecorder()
	tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(sr))
	ctx, root := tp.Tracer("test").Start(context.Background(), "root")
	fn(ctx)
	root.End()
	for _, s := range sr.Ended() {
		if s.Name() == "root" {
			m := make(map[attribute.Key]attribute.Value, len(s.Attributes()))
			for _, kv := range s.Attributes() {
				m[kv.Key] = kv.Value
			}
			return m
		}
	}
	t.Fatal("root span not recorded")
	return nil
}

func assertSpanStr(t *testing.T, attrs map[attribute.Key]attribute.Value, key attribute.Key, want string) {
	t.Helper()
	v, ok := attrs[key]
	if !ok {
		t.Errorf("missing %s", key)
		return
	}
	if v.AsString() != want {
		t.Errorf("%s = %q, want %q", key, v.AsString(), want)
	}
}

// foreignUID matches no stored object, for guards that must miss.
const foreignUID = "0f0e0d0c-0b0a-4908-8706-050403020100"