substrate/cmd/ateapi/internal/controlapi/functionaltest/main_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"
	"log"
	"net"
	"os"
	"strconv"
	"sync"
	"testing"
	"time"

	"go.ate.dev/substrate/cmd/ateapi/internal/store/storetest"
	"go.ate.dev/substrate/internal/atelet"
	"go.ate.dev/substrate/internal/objectstore/objectstoretest"
	"go.ate.dev/substrate/internal/proto/ateletpb"
	"go.ate.dev/substrate/internal/serverboot"
	"go.ate.dev/substrate/internal/testenv"
	"google.golang.org/grpc"
	"google.golang.org/protobuf/proto"
	corev1 "k8s.io/api/core/v1"
	storagev1 "k8s.io/api/storage/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
)

var (
	cfg        *rest.Config
	fakeAtelet = &FakeAteletServer{}
)

func TestMain(m *testing.M) {
	// The in-process server logs RPC bodies through ateinterceptors, which
	// relies on this handler to mask debug_redact fields.
	serverboot.InitLoggerWithWriter(os.Stderr)
	var stopEnv func()
	cfg, stopEnv = testenv.Start()

	// Create ate-system namespace
	k8sClient, err := kubernetes.NewForConfig(cfg)
	if err != nil {
		log.Fatalf("kubernetes.NewForConfig: %v", err)
	}
	_, err = k8sClient.CoreV1().Namespaces().Create(context.Background(), &corev1.Namespace{
		ObjectMeta: metav1.ObjectMeta{Name: "ate-system"},
	}, metav1.CreateOptions{})
	if err != nil && !apierrors.IsAlreadyExists(err) {
		log.Fatalf("create ate-system namespace: %v", err)
	}

	// Create StorageClasses for volume tests
	_, err = k8sClient.StorageV1().StorageClasses().Create(context.Background(), &storagev1.StorageClass{
		ObjectMeta:  metav1.ObjectMeta{Name: "standard"},
		Provisioner: "substrate.io/mock",
	}, metav1.CreateOptions{})
	if err != nil && !apierrors.IsAlreadyExists(err) {
		log.Fatalf("create standard storage class: %v", err)
	}
	_, err = k8sClient.StorageV1().StorageClasses().Create(context.Background(), &storagev1.StorageClass{
		ObjectMeta:  metav1.ObjectMeta{Name: "fast"},
		Provisioner: "substrate.io/mock",
	}, metav1.CreateOptions{})
	if err != nil && !apierrors.IsAlreadyExists(err) {
		log.Fatalf("create fast storage class: %v", err)
	}

	// Create shared Atelet Pod. Tests that place workers on another node add
	// their own atelet there via setupAteletOnNode.
	if err := createAteletPod(k8sClient, "atelet-shared", "node1"); err != nil {
		log.Fatalf("%v", err)
	}

	// Start Fake Atelet Server on the port ateapi dials.
	ateletGrpcServer := grpc.NewServer()
	ateletpb.RegisterAteletServer(ateletGrpcServer, fakeAtelet)
	ateletAddr := net.JoinHostPort("127.0.0.1", strconv.Itoa(atelet.DefaultPort))
	ateletLis, err := net.Listen("tcp", ateletAddr)
	if err != nil {
		log.Fatalf("listen on %s: %v", ateletAddr, err)
	}
	go func() {
		if err := ateletGrpcServer.Serve(ateletLis); err != nil {
			fmt.Printf("atelet grpc server exited: %v\n", err)
		}
	}()

	code := m.Run()

	ateletGrpcServer.Stop()

	stopEnv()

	storetest.Shutdown()

	os.Exit(code)
}

// FakeAteletServer implements ateletpb.WorkersServer
type FakeAteletServer struct {
	ateletpb.UnimplementedAteletServer

	Lock sync.Mutex

	RunCalled  bool
	RunRequest *ateletpb.RunRequest
	FailRun    error

	CheckpointCalled    bool
	CheckpointRequest   *ateletpb.CheckpointRequest
	FailCheckpoint      error
	KeepCheckpointLocal bool
	// checkpointed is each actor's last checkpointed snapshot, by UID.
	checkpointed map[string]*ateletpb.Snapshot

	RestoreCalled  bool
	RestoreRequest *ateletpb.RestoreRequest
	FailRestore    error
	RestoreDelay   time.Duration

	PromoteCalled  bool
	PromoteRequest *ateletpb.PromoteSnapshotRequest
	FailPromote    error

	DetachCalled  bool
	DetachRequest *ateletpb.DetachActorRequest
	FailDetach    error

	TerminateCalled  bool
	TerminateRequest *ateletpb.TerminateRequest
	FailTerminate    error

	// objectStore, when set, receives the objects a checkpoint or an upload
	// writes, so the control plane's copy and release steps have real external
	// snapshots to act on. setupTest points it at the test's own store.
	objectStore *objectstoretest.Fake
}

// snapshotObjects names the objects a checkpoint leaves behind. The names are
// arbitrary: nothing in the control plane reads them, it only copies and
// deletes whatever shares the snapshot's prefix.
var snapshotObjects = []string{"manifest.json", "memory.zst"}

// SetObjectStore points the fake at the store a checkpoint should write to.
func (f *FakeAteletServer) SetObjectStore(store *objectstoretest.Fake) {
	f.Lock.Lock()
	defer f.Lock.Unlock()
	f.objectStore = store
}

// writeSnapshot records the external snapshot a checkpoint or an upload would
// have written at snapshotURI. Called with Lock held.
func (f *FakeAteletServer) writeSnapshot(snapshotURI string) error {
	if f.objectStore == nil || snapshotURI == "" {
		return nil
	}
	return f.objectStore.WriteSnapshot(snapshotURI, snapshotObjects...)
}

func (f *FakeAteletServer) Reset() {
	f.Lock.Lock()
	defer f.Lock.Unlock()

	f.RunCalled = false
	f.RunRequest = nil
	f.FailRun = nil

	f.CheckpointCalled = false
	f.CheckpointRequest = nil
	f.FailCheckpoint = nil
	f.checkpointed = nil
	f.KeepCheckpointLocal = false

	f.RestoreCalled = false
	f.RestoreRequest = nil
	f.FailRestore = nil
	f.RestoreDelay = 0

	f.PromoteCalled = false
	f.PromoteRequest = nil
	f.FailPromote = nil

	f.DetachCalled = false
	f.DetachRequest = nil
	f.FailDetach = nil

	f.TerminateCalled = false
	f.TerminateRequest = nil
	f.FailTerminate = nil

	f.objectStore = nil
}

func (f *FakeAteletServer) PromoteSnapshot(ctx context.Context, req *ateletpb.PromoteSnapshotRequest) (*ateletpb.PromoteSnapshotResponse, error) {
	f.Lock.Lock()
	defer f.Lock.Unlock()

	f.PromoteCalled = true
	f.PromoteRequest = proto.Clone(req).(*ateletpb.PromoteSnapshotRequest)
	if f.FailPromote != nil {
		return nil, f.FailPromote
	}
	if err := f.writeSnapshot(req.GetSnapshot().GetObject().GetSnapshotUri()); err != nil {
		return nil, err
	}
	// Like atelet, report the snapshot durable and still kept on the node.
	report := &ateletpb.SnapshotReport{Snapshot: proto.CloneOf(req.GetSnapshot())}
	return &ateletpb.PromoteSnapshotResponse{State: &ateletpb.ActorSnapshotState{
		LatestLocal:   report,
		LatestDurable: proto.CloneOf(report),
	}}, nil
}

// DetachActor uploads the actor's last checkpointed snapshot and reports it
// durable, with nothing left on the node.
func (f *FakeAteletServer) DetachActor(ctx context.Context, req *ateletpb.DetachActorRequest) (*ateletpb.DetachActorResponse, error) {
	f.Lock.Lock()
	defer f.Lock.Unlock()

	f.DetachCalled = true
	f.DetachRequest = proto.CloneOf(req)
	if f.FailDetach != nil {
		return nil, f.FailDetach
	}
	snap := f.checkpointed[req.GetActorUid()]
	if snap == nil {
		return &ateletpb.DetachActorResponse{State: &ateletpb.ActorSnapshotState{}}, nil
	}
	if err := f.writeSnapshot(snap.GetObject().GetSnapshotUri()); err != nil {
		return nil, err
	}
	return &ateletpb.DetachActorResponse{State: &ateletpb.ActorSnapshotState{
		LatestDurable: &ateletpb.SnapshotReport{Snapshot: proto.CloneOf(snap)},
	}}, nil
}

func (f *FakeAteletServer) Run(ctx context.Context, req *ateletpb.RunRequest) (*ateletpb.RunResponse, error) {
	f.Lock.Lock()
	defer f.Lock.Unlock()

	f.RunCalled = true
	f.RunRequest = proto.Clone(req).(*ateletpb.RunRequest)
	if f.FailRun != nil {
		return nil, f.FailRun
	}

	return &ateletpb.RunResponse{}, nil
}

func (f *FakeAteletServer) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRequest) (*ateletpb.CheckpointResponse, error) {
	f.Lock.Lock()
	defer f.Lock.Unlock()

	f.CheckpointCalled = true
	f.CheckpointRequest = proto.Clone(req).(*ateletpb.CheckpointRequest)
	if f.FailCheckpoint != nil {
		return nil, f.FailCheckpoint
	}
	if f.checkpointed == nil {
		f.checkpointed = map[string]*ateletpb.Snapshot{}
	}
	f.checkpointed[req.GetActorUid()] = proto.CloneOf(req.GetSnapshot())

	// Only an immediate upload is in object storage when Checkpoint returns;
	// atelet does not upload deferred snapshots yet.
	if req.GetUploadConfig().GetMode() == ateletpb.UploadMode_UPLOAD_MODE_IMMEDIATE {
		if err := f.writeSnapshot(req.GetSnapshot().GetObject().GetSnapshotUri()); err != nil {
			return nil, err
		}
	}
	if f.KeepCheckpointLocal {
		report := &ateletpb.SnapshotReport{Snapshot: proto.CloneOf(req.GetSnapshot())}
		state := &ateletpb.ActorSnapshotState{LatestLocal: report}
		if req.GetUploadConfig().GetMode() == ateletpb.UploadMode_UPLOAD_MODE_IMMEDIATE {
			state.LatestDurable = proto.CloneOf(report)
		}
		return &ateletpb.CheckpointResponse{State: state}, nil
	}
	return &ateletpb.CheckpointResponse{}, nil
}

func (f *FakeAteletServer) Restore(ctx context.Context, req *ateletpb.RestoreRequest) (*ateletpb.RestoreResponse, error) {
	f.Lock.Lock()
	defer f.Lock.Unlock()

	f.RestoreCalled = true
	f.RestoreRequest = proto.Clone(req).(*ateletpb.RestoreRequest)
	if f.RestoreDelay > 0 {
		time.Sleep(f.RestoreDelay)
	}
	if f.FailRestore != nil {
		return nil, f.FailRestore
	}
	return &ateletpb.RestoreResponse{}, nil
}

func (f *FakeAteletServer) Terminate(ctx context.Context, req *ateletpb.TerminateRequest) (*ateletpb.TerminateResponse, error) {
	f.Lock.Lock()
	defer f.Lock.Unlock()

	f.TerminateCalled = true
	f.TerminateRequest = proto.Clone(req).(*ateletpb.TerminateRequest)
	if f.FailTerminate != nil {
		return nil, f.FailTerminate
	}
	return &ateletpb.TerminateResponse{}, nil
}

func (f *FakeAteletServer) lastRestoreRequest() *ateletpb.RestoreRequest {
	f.Lock.Lock()
	defer f.Lock.Unlock()

	if f.RestoreRequest == nil {
		return nil
	}
	return proto.Clone(f.RestoreRequest).(*ateletpb.RestoreRequest)
}

func (f *FakeAteletServer) lastRestoreVolumePublishContext(volumeName string) map[string]string {
	for _, v := range f.lastRestoreRequest().GetSpec().GetVolumes() {
		if v.GetName() == volumeName {
			return v.GetExternal().GetPublishContext()
		}
	}
	return nil
}