atelet: add successful RPC boundary tests

cmd/atelet/main.go

@@ -67,7 +67,7 @@ var (
 
 	gcpAuthForImagePulls         = pflag.Bool("gcp-auth-for-image-pulls", true, "Use GCP application default credentials mechanism.")
 	localhostRegistryReplacement = pflag.String("localhost-registry-replacement", "", "The replacement registry endpoint for localhost and/or loopback IP addresses, useful for local development. for example kind-registry:5000")
-	imageCacheDir                = pflag.String("image-cache-dir", ateompath.ImageCacheDir, "Directory for the node-local OCI image layer cache. Must be on the volume shared with the ateom pods (the cached layers are their overlay lowerdirs), and on a disk sized for both capacity and IOPS: unpack throughput is gated by the volume's IOPS.")
+	imageCacheDir                = pflag.String("image-cache-dir", ateompath.ImageCacheDir(), "Directory for the node-local OCI image layer cache. Must be on the volume shared with the ateom pods (the cached layers are their overlay lowerdirs), and on a disk sized for both capacity and IOPS: unpack throughput is gated by the volume's IOPS.")
 
 	showVersion = pflag.Bool("version", false, "Print version and exit.")
 )
@@ -191,21 +191,25 @@ func main() {
 type AteomHerder struct {
 	ateletpb.UnimplementedAteomHerderServer
 
-	ateomDialer   *AteomDialer
-	imageCache    *imagecache.Store
+	ateomDialer   ateomClientDialer
+	imageCache    imageEnsurer
 	anonGCSClient ategcs.ObjectStorage
 	gcsClient     ategcs.ObjectStorage
 }
 
+type ateomClientDialer interface {
+	DialAteom(context.Context, string) (ateompb.AteomClient, error)
+}
+
 var _ ateletpb.AteomHerderServer = (*AteomHerder)(nil)
 
 // NewService creates a new WorkersManagerService.
 func NewService(
 	ctx context.Context,
-	ateomDialer *AteomDialer,
+	ateomDialer ateomClientDialer,
 	anonGCSClient ategcs.ObjectStorage,
 	gcsClient ategcs.ObjectStorage,
-	imageCache *imagecache.Store,
+	imageCache imageEnsurer,
 ) *AteomHerder {
 	wms := &AteomHerder{
 		ateomDialer:   ateomDialer,
@@ -759,11 +763,11 @@ func (s *AteomHerder) prepareOCIBundles(
 // dialAteom opens (or reuses) the gRPC connection to the target ateom
 // pod and returns an ateom client.
 func (s *AteomHerder) dialAteom(ctx context.Context, targetAteomUid string) (ateompb.AteomClient, error) {
-	conn, err := s.ateomDialer.DialAteomPod(ctx, targetAteomUid)
+	client, err := s.ateomDialer.DialAteom(ctx, targetAteomUid)
 	if err != nil {
-		return nil, fmt.Errorf("while getting ateom conn for %s: %w", targetAteomUid, err)
+		return nil, fmt.Errorf("while getting ateom client for %s: %w", targetAteomUid, err)
 	}
-	return ateompb.NewAteomClient(conn), nil
+	return client, nil
 }
 
 // buildAteomWorkloadSpec projects the atelet-facing workload spec onto
@@ -814,6 +818,14 @@ type AteomDialer struct {
 	conns *lru.Cache
 }
 
+func (d *AteomDialer) DialAteom(ctx context.Context, podUID string) (ateompb.AteomClient, error) {
+	conn, err := d.DialAteomPod(ctx, podUID)
+	if err != nil {
+		return nil, err
+	}
+	return ateompb.NewAteomClient(conn), nil
+}
+
 func (d *AteomDialer) DialAteomPod(ctx context.Context, podUID string) (*grpc.ClientConn, error) {
 	key := podUID
 

cmd/atelet/main_test.go

@@ -23,20 +23,36 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"runtime"
 	"strings"
 	"syscall"
 	"testing"
 
 	"github.com/agent-substrate/substrate/internal/ateerrors"
 	"github.com/agent-substrate/substrate/internal/ateompath"
+	"github.com/agent-substrate/substrate/internal/imagecache"
 	"github.com/agent-substrate/substrate/internal/proto/ateletpb"
 	"github.com/agent-substrate/substrate/internal/proto/ateompb"
 	"github.com/google/go-cmp/cmp"
+	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/testing/protocmp"
 )
 
+// setAteomBasePathForTest mutates shared package state and must not be used by
+// tests that run in parallel.
+func setAteomBasePathForTest(t *testing.T) {
+	t.Helper()
+	orig := ateompath.BasePath
+	ateompath.BasePath = t.TempDir()
+	t.Cleanup(func() { ateompath.BasePath = orig })
+	if err := os.MkdirAll(ateompath.StaticFilesDir(), 0o700); err != nil {
+		t.Fatalf("creating static files dir: %v", err)
+	}
+}
+
 func TestWriteFileAtomic(t *testing.T) {
 	dir := t.TempDir()
 	target := filepath.Join(dir, "actor-id")
@@ -260,15 +276,13 @@ func TestValidateRestoreRequest(t *testing.T) {
 // prove the ordering, it plants a real file at the exact path an invalid hash
 // resolves to: a correctly-ordered fetchAsset validates first and returns an
 // error, while a regression that stats first would find this file and return it
-// with a nil error, failing the test. StaticFilesDir is redirected to a temp
-// dir so the planted path is writable and isolated.
+// with a nil error, failing the test. BasePath is redirected to a temp dir so
+// the derived static-files path is writable and isolated.
 func TestFetchAssetRejectsBadHash(t *testing.T) {
-	orig := ateompath.StaticFilesDir
-	ateompath.StaticFilesDir = t.TempDir()
-	t.Cleanup(func() { ateompath.StaticFilesDir = orig })
+	setAteomBasePathForTest(t)
 
 	// Invalid (8 chars, not 64) but separator-free, so it resolves to a normal
-	// filename inside the temp StaticFilesDir.
+	// filename inside the temp static-files directory.
 	const badHash = "deadbeef"
 	if err := os.WriteFile(ateompath.RunSCBinaryPath(badHash), []byte("planted"), 0o755); err != nil {
 		t.Fatalf("planting cache file: %v", err)
@@ -304,15 +318,15 @@ func (fakeObjectStorage) PutObject(_ context.Context, _, _ string, _ io.Reader)
 // TestFetchAssetStreaming covers the streamed download: good asset cached,
 // over-cap rejected, hash mismatch rejected (failures leave no cache file).
 func TestFetchAssetStreaming(t *testing.T) {
-	origDir, origCap := ateompath.StaticFilesDir, maxAssetBytes
-	t.Cleanup(func() { ateompath.StaticFilesDir, maxAssetBytes = origDir, origCap })
+	origCap := maxAssetBytes
+	t.Cleanup(func() { maxAssetBytes = origCap })
 
 	content := []byte("micro-vm kernel bytes")
 	goodHash := fmt.Sprintf("%x", sha256.Sum256(content))
 	const url = "gs://test-bucket/asset"
 
 	t.Run("good asset is cached", func(t *testing.T) {
-		ateompath.StaticFilesDir = t.TempDir()
+		setAteomBasePathForTest(t)
 		s := &AteomHerder{anonGCSClient: fakeObjectStorage{data: content}}
 		path, err := s.fetchAsset(context.Background(), assetEntry{URL: url, SHA256: goodHash})
 		if err != nil {
@@ -328,7 +342,7 @@ func TestFetchAssetStreaming(t *testing.T) {
 	})
 
 	t.Run("over-cap asset rejected, cache not written", func(t *testing.T) {
-		ateompath.StaticFilesDir = t.TempDir()
+		setAteomBasePathForTest(t)
 		maxAssetBytes = 4 // content is longer than this
 		s := &AteomHerder{anonGCSClient: fakeObjectStorage{data: content}}
 		_, err := s.fetchAsset(context.Background(), assetEntry{URL: url, SHA256: goodHash})
@@ -344,7 +358,7 @@ func TestFetchAssetStreaming(t *testing.T) {
 	})
 
 	t.Run("hash mismatch rejected, cache not written", func(t *testing.T) {
-		ateompath.StaticFilesDir = t.TempDir()
+		setAteomBasePathForTest(t)
 		maxAssetBytes = origCap
 		wrongHash := strings.Repeat("a", 64) // valid 64-hex format, wrong value
 		s := &AteomHerder{anonGCSClient: fakeObjectStorage{data: content}}
@@ -361,7 +375,7 @@ func TestFetchAssetStreaming(t *testing.T) {
 	})
 
 	t.Run("missing object is terminal", func(t *testing.T) {
-		ateompath.StaticFilesDir = t.TempDir()
+		setAteomBasePathForTest(t)
 		maxAssetBytes = origCap
 		// The ategcs clients tag a missing object with ReasonFailedGetExternalObject.
 		notFound := fmt.Errorf("%w: no such object", ateerrors.ReasonFailedGetExternalObject)
@@ -381,7 +395,7 @@ func TestFetchAssetStreaming(t *testing.T) {
 	})
 
 	t.Run("malformed url is terminal", func(t *testing.T) {
-		ateompath.StaticFilesDir = t.TempDir()
+		setAteomBasePathForTest(t)
 		maxAssetBytes = origCap
 		s := &AteomHerder{anonGCSClient: fakeObjectStorage{data: content}}
 		// Invalid percent-escape: url.Parse rejects it inside ategcs.Open, which
@@ -393,7 +407,7 @@ func TestFetchAssetStreaming(t *testing.T) {
 	})
 
 	t.Run("network error stays untagged (retriable)", func(t *testing.T) {
-		ateompath.StaticFilesDir = t.TempDir()
+		setAteomBasePathForTest(t)
 		maxAssetBytes = origCap
 		s := &AteomHerder{anonGCSClient: fakeObjectStorage{err: errors.New("connection refused")}}
 		_, err := s.fetchAsset(context.Background(), assetEntry{URL: url, SHA256: goodHash})
@@ -458,6 +472,154 @@ func TestRPCBoundariesReject(t *testing.T) {
 	})
 }
 
+type fakeImageEnsurer struct{}
+
+func (fakeImageEnsurer) EnsureImage(context.Context, string) (*imagecache.Image, error) {
+	return &imagecache.Image{Config: v1.Config{}}, nil
+}
+
+type fakeAteomClient struct {
+	runReq        *ateompb.RunWorkloadRequest
+	checkpointReq *ateompb.CheckpointWorkloadRequest
+	restoreReq    *ateompb.RestoreWorkloadRequest
+}
+
+func (c *fakeAteomClient) RunWorkload(_ context.Context, req *ateompb.RunWorkloadRequest, _ ...grpc.CallOption) (*ateompb.RunWorkloadResponse, error) {
+	c.runReq = req
+	return &ateompb.RunWorkloadResponse{}, nil
+}
+
+func (c *fakeAteomClient) CheckpointWorkload(_ context.Context, req *ateompb.CheckpointWorkloadRequest, _ ...grpc.CallOption) (*ateompb.CheckpointWorkloadResponse, error) {
+	c.checkpointReq = req
+	const snapshotFile = "checkpoint.img"
+	if err := os.WriteFile(filepath.Join(ateompath.CheckpointStateDir(req.GetActorUid()), snapshotFile), []byte("snapshot"), 0o600); err != nil {
+		return nil, fmt.Errorf("writing fake checkpoint: %w", err)
+	}
+	return &ateompb.CheckpointWorkloadResponse{SnapshotFiles: []string{snapshotFile}}, nil
+}
+
+func (c *fakeAteomClient) RestoreWorkload(_ context.Context, req *ateompb.RestoreWorkloadRequest, _ ...grpc.CallOption) (*ateompb.RestoreWorkloadResponse, error) {
+	c.restoreReq = req
+	return &ateompb.RestoreWorkloadResponse{}, nil
+}
+
+type fakeAteomDialer struct {
+	client  ateompb.AteomClient
+	targets []string
+}
+
+func (d *fakeAteomDialer) DialAteom(_ context.Context, target string) (ateompb.AteomClient, error) {
+	d.targets = append(d.targets, target)
+	return d.client, nil
+}
+
+// TestRPCBoundariesAccept runs a successful local lifecycle through the public
+// atelet RPC methods. The fake dependencies stop at the image-fetch and ateom
+// client boundaries; filesystem setup, request projection, checkpoint movement,
+// and sandbox record handling all use the production paths.
+func TestRPCBoundariesAccept(t *testing.T) {
+	// Do not call t.Parallel: the test redirects the package-wide BasePath.
+	setAteomBasePathForTest(t)
+
+	if err := os.MkdirAll(ateompath.StaticFilesDir(), 0o700); err != nil {
+		t.Fatalf("creating static files dir: %v", err)
+	}
+	runscSum := fmt.Sprintf("%x", sha256.Sum256([]byte("fake runsc")))
+	runscPath := ateompath.RunSCBinaryPath(runscSum)
+	if err := os.WriteFile(runscPath, []byte("fake runsc"), 0o755); err != nil {
+		t.Fatalf("seeding runsc cache: %v", err)
+	}
+
+	client := &fakeAteomClient{}
+	dialer := &fakeAteomDialer{client: client}
+	// GCS clients stay nil: the runsc asset is pre-seeded above, so
+	// ensureSandboxAssets never downloads, and checkpoint/restore use local
+	// snapshots, so no path touches object storage.
+	s := &AteomHerder{ateomDialer: dialer, imageCache: fakeImageEnsurer{}}
+	ctx := context.Background()
+	const snapshotPrefix = "snapshot-1"
+
+	runReq := validRunRequest()
+	runReq.Spec.PauseImage = "pause:test"
+	runReq.Spec.Containers[0].Image = "worker:test"
+	runReq.Spec.Containers[0].Command = []string{"/worker"}
+	runReq.SandboxAssets = &ateletpb.SandboxAssets{
+		SandboxClass: "gvisor",
+		Assets: map[string]*ateletpb.ArchAssets{
+			runtime.GOARCH: {
+				Files: map[string]*ateletpb.AssetFile{
+					"runsc": {Url: "gs://unused/runsc", Sha256: runscSum},
+				},
+			},
+		},
+	}
+
+	t.Run("Run", func(t *testing.T) {
+		if _, err := s.Run(ctx, runReq); err != nil {
+			t.Fatalf("Run: %v", err)
+		}
+		if client.runReq == nil {
+			t.Fatal("Run did not call ateom")
+		}
+		if got := client.runReq.GetRunscPath(); got != runscPath {
+			t.Errorf("runsc path = %q, want %q", got, runscPath)
+		}
+		if got := client.runReq.GetActorUid(); got != runReq.GetActorUid() {
+			t.Errorf("actor UID = %q, want %q", got, runReq.GetActorUid())
+		}
+	})
+
+	checkpointReq := validCheckpointRequest()
+	checkpointReq.Spec = runReq.Spec
+	checkpointReq.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL
+	checkpointReq.Config = &ateletpb.CheckpointRequest_LocalConfig{
+		LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: snapshotPrefix},
+	}
+	t.Run("Checkpoint", func(t *testing.T) {
+		if _, err := s.Checkpoint(ctx, checkpointReq); err != nil {
+			t.Fatalf("Checkpoint: %v", err)
+		}
+		if client.checkpointReq == nil {
+			t.Fatal("Checkpoint did not call ateom")
+		}
+		if got := client.checkpointReq.GetRunscPath(); got != runscPath {
+			t.Errorf("runsc path = %q, want %q", got, runscPath)
+		}
+		localSnapshot := filepath.Join(ateompath.LocalCheckpointsDir(checkpointReq.GetActorUid()), snapshotPrefix, "checkpoint.img")
+		if got, err := os.ReadFile(localSnapshot); err != nil {
+			t.Errorf("reading local snapshot: %v", err)
+		} else if string(got) != "snapshot" {
+			t.Errorf("local snapshot = %q, want %q", got, "snapshot")
+		}
+	})
+
+	restoreReq := validRestoreRequest()
+	restoreReq.Spec = runReq.Spec
+	restoreReq.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL
+	restoreReq.Config = &ateletpb.RestoreRequest_LocalConfig{
+		LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: snapshotPrefix},
+	}
+	t.Run("Restore", func(t *testing.T) {
+		if _, err := s.Restore(ctx, restoreReq); err != nil {
+			t.Fatalf("Restore: %v", err)
+		}
+		if client.restoreReq == nil {
+			t.Fatal("Restore did not call ateom")
+		}
+		if got := client.restoreReq.GetRunscPath(); got != runscPath {
+			t.Errorf("runsc path = %q, want %q", got, runscPath)
+		}
+		if got := client.restoreReq.GetActorUid(); got != restoreReq.GetActorUid() {
+			t.Errorf("actor UID = %q, want %q", got, restoreReq.GetActorUid())
+		}
+	})
+
+	wantTargets := []string{runReq.GetTargetAteomUid(), checkpointReq.GetTargetAteomUid(), restoreReq.GetTargetAteomUid()}
+	if diff := cmp.Diff(wantTargets, dialer.targets); diff != "" {
+		t.Errorf("dial targets mismatch (-want +got):\n%s", diff)
+	}
+}
+
 func TestBuildAteomWorkloadSpecForwardsReadyz(t *testing.T) {
 	in := &ateletpb.WorkloadSpec{
 		PauseImage: "pause",

cmd/atelet/oci.go

@@ -50,7 +50,11 @@ const (
 	ActorIDFileName = "actor-id"
 )
 
-func prepareOCIDirectory(ctx context.Context, imageCache *imagecache.Store, actorUID, containerName, ref string, command, args []string, env []string, annotations map[string]string, netns string, identityDir string, durableDirVolumeMounts []*ateletpb.VolumeMount) error {
+type imageEnsurer interface {
+	EnsureImage(context.Context, string) (*imagecache.Image, error)
+}
+
+func prepareOCIDirectory(ctx context.Context, imageCache imageEnsurer, actorUID, containerName, ref string, command, args []string, env []string, annotations map[string]string, netns string, identityDir string, durableDirVolumeMounts []*ateletpb.VolumeMount) error {
 	tracer := otel.Tracer("prepareOCIDirectory")
 
 	ctx, span := tracer.Start(ctx, "prepareOCIDirectory")

cmd/atelet/sandbox_assets.go

@@ -93,7 +93,7 @@ func recordFromRequest(sa *ateletpb.SandboxAssets) (*sandboxAssetsRecord, error)
 // the micro-VM runtime has several (kata-shim, cloud-hypervisor, ...). Assets are
 // cached, so re-fetching at Checkpoint/Restore is a no-op once present.
 func (s *AteomHerder) ensureSandboxAssets(ctx context.Context, rec *sandboxAssetsRecord) (map[string]string, error) {
-	if err := os.MkdirAll(ateompath.StaticFilesDir, 0o700); err != nil {
+	if err := os.MkdirAll(ateompath.StaticFilesDir(), 0o700); err != nil {
 		if isTerminalFileSystemErr(err) {
 			return nil, fmt.Errorf("%w: while creating static files dir: %w", ateerrors.ReasonTerminalFileSystemError, err)
 		}