refactor: move ami selection to aws foundation

This commit is contained in:
joshsiv-mulga
2026-10-02 13:39:41 +10:00
parent 30c7541a2f
commit b5bc8cd7d6
8 changed files with 148 additions and 134 deletions
+40
View File
@@ -0,0 +1,40 @@
// Package ami provides small AWS AMI model helpers shared by service domains.
package ami
import (
"slices"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ec2"
)
// SelectNewestImage returns the newest image ID among images, its CreationDate,
// and how many images were considered. Nil images, images with an empty ImageId
// and nil tag entries are skipped; a non-empty excludeTagKey drops images
// carrying that tag key. Returns an empty ID when nothing matches.
func SelectNewestImage(images []*ec2.Image, excludeTagKey string) (imageID, created string, matches int) {
for _, img := range images {
if img == nil || aws.StringValue(img.ImageId) == "" {
continue
}
// DescribeImages filters have no negation, so callers that must reject a
// tag (e.g. gpu-vendor) can only do it client-side, here.
if excludeTagKey != "" && imageHasTagKey(img, excludeTagKey) {
continue
}
matches++
// CreationDate is a fixed-width RFC3339 timestamp, so lexicographic
// comparison orders it correctly without parsing.
if c := aws.StringValue(img.CreationDate); imageID == "" || c > created {
imageID, created = aws.StringValue(img.ImageId), c
}
}
return imageID, created, matches
}
// imageHasTagKey reports whether img carries a tag with the given key.
func imageHasTagKey(img *ec2.Image, key string) bool {
return slices.ContainsFunc(img.Tags, func(t *ec2.Tag) bool {
return t != nil && aws.StringValue(t.Key) == key
})
}
@@ -0,0 +1,97 @@
package ami
import (
"testing"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ec2"
)
func imageForTest(id, created string, tagKeys ...string) *ec2.Image {
image := &ec2.Image{ImageId: aws.String(id), CreationDate: aws.String(created)}
for _, key := range tagKeys {
if key == "" {
image.Tags = append(image.Tags, nil)
continue
}
image.Tags = append(image.Tags, &ec2.Tag{Key: aws.String(key), Value: aws.String("v")})
}
return image
}
func TestSelectNewestImage(t *testing.T) {
tests := []struct {
name string
images []*ec2.Image
excludeKey string
wantID string
wantCreated string
wantMatches int
}{
{
name: "empty input",
},
{
name: "all entries nil",
images: []*ec2.Image{nil, nil},
},
{
name: "empty image id skipped",
images: []*ec2.Image{imageForTest("", "2026-01-01T00:00:00.000Z"), imageForTest("ami-a", "2025-01-01T00:00:00.000Z")},
// The empty-ID entry is not a match, so no multi-match warning fires.
wantID: "ami-a", wantCreated: "2025-01-01T00:00:00.000Z", wantMatches: 1,
},
{
name: "newest creation date wins regardless of order",
images: []*ec2.Image{imageForTest("ami-old", "2026-01-01T00:00:00.000Z"), imageForTest("ami-new", "2026-08-01T00:00:00.000Z"), imageForTest("ami-mid", "2026-04-01T00:00:00.000Z")},
wantID: "ami-new", wantCreated: "2026-08-01T00:00:00.000Z", wantMatches: 3,
},
{
name: "identical dates keep the first match",
images: []*ec2.Image{imageForTest("ami-a", "2026-01-01T00:00:00.000Z"), imageForTest("ami-b", "2026-01-01T00:00:00.000Z")},
wantID: "ami-a", wantCreated: "2026-01-01T00:00:00.000Z", wantMatches: 2,
},
{
name: "exclude tag key drops tagged images",
images: []*ec2.Image{imageForTest("ami-gpu", "2026-08-01T00:00:00.000Z", "gpu-vendor"), imageForTest("ami-cpu", "2026-01-01T00:00:00.000Z")},
excludeKey: "gpu-vendor",
wantID: "ami-cpu", wantCreated: "2026-01-01T00:00:00.000Z", wantMatches: 1,
},
{
name: "empty exclude key keeps tagged images",
images: []*ec2.Image{imageForTest("ami-gpu", "2026-08-01T00:00:00.000Z", "gpu-vendor"), imageForTest("ami-cpu", "2026-01-01T00:00:00.000Z")},
wantID: "ami-gpu", wantCreated: "2026-08-01T00:00:00.000Z", wantMatches: 2,
},
{
name: "all images excluded",
images: []*ec2.Image{imageForTest("ami-gpu", "2026-08-01T00:00:00.000Z", "gpu-vendor")},
excludeKey: "gpu-vendor",
},
{
name: "nil tag entry does not panic",
images: []*ec2.Image{imageForTest("ami-a", "2026-01-01T00:00:00.000Z", "", "gpu-vendor"), imageForTest("ami-b", "2026-02-01T00:00:00.000Z", "")},
excludeKey: "gpu-vendor",
wantID: "ami-b", wantCreated: "2026-02-01T00:00:00.000Z", wantMatches: 1,
},
{
name: "missing creation date sorts oldest",
images: []*ec2.Image{{ImageId: aws.String("ami-nodate")}, imageForTest("ami-dated", "2020-01-01T00:00:00.000Z")},
wantID: "ami-dated", wantCreated: "2020-01-01T00:00:00.000Z", wantMatches: 2,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
gotID, gotCreated, gotMatches := SelectNewestImage(tt.images, tt.excludeKey)
if gotID != tt.wantID {
t.Errorf("imageID = %q, want %q", gotID, tt.wantID)
}
if gotCreated != tt.wantCreated {
t.Errorf("created = %q, want %q", gotCreated, tt.wantCreated)
}
if gotMatches != tt.wantMatches {
t.Errorf("matches = %d, want %d", gotMatches, tt.wantMatches)
}
})
}
}
+2 -1
View File
@@ -8,6 +8,7 @@ import (
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ec2"
awsami "github.com/mulgadc/spinifex/spinifex/foundation/aws/ami"
"github.com/mulgadc/spinifex/spinifex/foundation/aws/tags"
"github.com/mulgadc/spinifex/spinifex/utils"
)
@@ -44,7 +45,7 @@ func resolveServingAMI(ctx context.Context, imgSvc amiResolver) (string, error)
}
// Serving images are GPU-tagged by design, so nothing is excluded here.
newestID, _, matches := utils.SelectNewestImage(out.Images, "")
newestID, _, matches := awsami.SelectNewestImage(out.Images, "")
if newestID == "" {
return "", ErrServingAMINotFound
}
+4 -4
View File
@@ -8,11 +8,11 @@ import (
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ec2"
"github.com/mulgadc/spinifex/spinifex/foundation/aws/errors"
"github.com/mulgadc/spinifex/spinifex/domains/ec2/instancetypes"
"github.com/mulgadc/spinifex/spinifex/domains/ec2/ebs/metadata"
"github.com/mulgadc/spinifex/spinifex/domains/ec2/instancetypes"
awsami "github.com/mulgadc/spinifex/spinifex/foundation/aws/ami"
"github.com/mulgadc/spinifex/spinifex/foundation/aws/errors"
"github.com/mulgadc/spinifex/spinifex/foundation/aws/tags"
"github.com/mulgadc/spinifex/spinifex/utils"
)
const (
@@ -212,7 +212,7 @@ func resolveNewestAMI(ctx context.Context, amiSvc ecsImageResolver, accountID st
if excludeGPUTagged {
excludeKey = tags.GPUVendorKey
}
newestID, newestCreated, matches := utils.SelectNewestImage(out.Images, excludeKey)
newestID, newestCreated, matches := awsami.SelectNewestImage(out.Images, excludeKey)
if newestID == "" {
return "", fmt.Errorf("%w (%s, account %s)", notFound, notFoundDesc, accountID)
}
+3 -3
View File
@@ -13,10 +13,10 @@ import (
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ec2"
awsami "github.com/mulgadc/spinifex/spinifex/foundation/aws/ami"
"github.com/mulgadc/spinifex/spinifex/foundation/aws/errors"
"github.com/mulgadc/spinifex/spinifex/handlers/sysinstance"
"github.com/mulgadc/spinifex/spinifex/foundation/aws/tags"
"github.com/mulgadc/spinifex/spinifex/utils"
"github.com/mulgadc/spinifex/spinifex/handlers/sysinstance"
)
// ErrEKSServerAMINotFound is returned when no AMI with the EKS managed-by tag
@@ -393,7 +393,7 @@ func resolveNewestAMI(ctx context.Context, amiSvc k3sAMIResolver, accountID stri
if excludeGPUTagged {
excludeKey = tags.GPUVendorKey
}
newestID, newestCreated, matches := utils.SelectNewestImage(out.Images, excludeKey)
newestID, newestCreated, matches := awsami.SelectNewestImage(out.Images, excludeKey)
if newestID == "" {
return "", fmt.Errorf("%w (%s, account %s)", notFound, notFoundDesc, accountID)
}
+2 -1
View File
@@ -13,6 +13,7 @@ import (
"github.com/aws/aws-sdk-go/service/ec2"
"github.com/mulgadc/spinifex/contracts/ec2/v1"
"github.com/mulgadc/spinifex/spinifex/bootstrap/config"
awsami "github.com/mulgadc/spinifex/spinifex/foundation/aws/ami"
"github.com/mulgadc/spinifex/spinifex/foundation/aws/errors"
"github.com/mulgadc/spinifex/spinifex/foundation/aws/tags"
"github.com/mulgadc/spinifex/spinifex/handlers/sysinstance"
@@ -542,7 +543,7 @@ func resolveEngineAMI(ctx context.Context, amiSvc launchAMIResolver, engine, ver
// Several builds of one engine version can be registered; select the most
// recently imported usable image. A GPU engine build carries the same engine
// tags, so it is excluded or a newer GPU image would hijack an ordinary instance.
newestID, _, matches := utils.SelectNewestImage(out.Images, tags.GPUVendorKey)
newestID, _, matches := awsami.SelectNewestImage(out.Images, tags.GPUVendorKey)
if newestID == "" {
return "", engineAMINotFound(engine, version)
}
-34
View File
@@ -24,12 +24,9 @@ import (
"os/signal"
"path/filepath"
"runtime"
"slices"
"strings"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ec2"
"github.com/klauspost/compress/zstd"
)
@@ -324,37 +321,6 @@ func hashImageFile(imagePath string, hasher hash.Hash) (string, error) {
return hex.EncodeToString(hasher.Sum(nil)), nil
}
// SelectNewestImage returns the newest image ID among images, its CreationDate,
// and how many images were considered. Nil images, images with an empty ImageId
// and nil tag entries are skipped; a non-empty excludeTagKey drops images
// carrying that tag key. Returns an empty ID when nothing matches.
func SelectNewestImage(images []*ec2.Image, excludeTagKey string) (imageID, created string, matches int) {
for _, img := range images {
if img == nil || aws.StringValue(img.ImageId) == "" {
continue
}
// DescribeImages filters have no negation, so callers that must reject a
// tag (e.g. gpu-vendor) can only do it client-side, here.
if excludeTagKey != "" && imageHasTagKey(img, excludeTagKey) {
continue
}
matches++
// CreationDate is a fixed-width RFC3339 timestamp, so lexicographic
// comparison orders it correctly without parsing.
if c := aws.StringValue(img.CreationDate); imageID == "" || c > created {
imageID, created = aws.StringValue(img.ImageId), c
}
}
return imageID, created, matches
}
// imageHasTagKey reports whether img carries a tag with the given key.
func imageHasTagKey(img *ec2.Image, key string) bool {
return slices.ContainsFunc(img.Tags, func(t *ec2.Tag) bool {
return t != nil && aws.StringValue(t.Key) == key
})
}
type Images struct {
Name string `json:"name"`
Description string `json:"description"`
-91
View File
@@ -11,8 +11,6 @@ import (
"syscall"
"testing"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ec2"
"github.com/klauspost/compress/zstd"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -61,95 +59,6 @@ func TestWithheldImages_VLLMServingCarriesBedrockTags(t *testing.T) {
}
}
func img(id, created string, tagKeys ...string) *ec2.Image {
image := &ec2.Image{ImageId: aws.String(id), CreationDate: aws.String(created)}
for _, k := range tagKeys {
if k == "" {
image.Tags = append(image.Tags, nil)
continue
}
image.Tags = append(image.Tags, &ec2.Tag{Key: aws.String(k), Value: aws.String("v")})
}
return image
}
func TestSelectNewestImage(t *testing.T) {
tests := []struct {
name string
images []*ec2.Image
excludeKey string
wantID string
wantCreated string
wantMatches int
}{
{
name: "empty input",
},
{
name: "all entries nil",
images: []*ec2.Image{nil, nil},
},
{
name: "empty image id skipped",
images: []*ec2.Image{img("", "2026-01-01T00:00:00.000Z"), img("ami-a", "2025-01-01T00:00:00.000Z")},
// The empty-ID entry is not a match, so no multi-match warning fires.
wantID: "ami-a", wantCreated: "2025-01-01T00:00:00.000Z", wantMatches: 1,
},
{
name: "newest creation date wins regardless of order",
images: []*ec2.Image{img("ami-old", "2026-01-01T00:00:00.000Z"), img("ami-new", "2026-08-01T00:00:00.000Z"), img("ami-mid", "2026-04-01T00:00:00.000Z")},
wantID: "ami-new", wantCreated: "2026-08-01T00:00:00.000Z", wantMatches: 3,
},
{
name: "identical dates keep the first match",
images: []*ec2.Image{img("ami-a", "2026-01-01T00:00:00.000Z"), img("ami-b", "2026-01-01T00:00:00.000Z")},
wantID: "ami-a", wantCreated: "2026-01-01T00:00:00.000Z", wantMatches: 2,
},
{
name: "exclude tag key drops tagged images",
images: []*ec2.Image{img("ami-gpu", "2026-08-01T00:00:00.000Z", "gpu-vendor"), img("ami-cpu", "2026-01-01T00:00:00.000Z")},
excludeKey: "gpu-vendor",
wantID: "ami-cpu", wantCreated: "2026-01-01T00:00:00.000Z", wantMatches: 1,
},
{
name: "empty exclude key keeps tagged images",
images: []*ec2.Image{img("ami-gpu", "2026-08-01T00:00:00.000Z", "gpu-vendor"), img("ami-cpu", "2026-01-01T00:00:00.000Z")},
wantID: "ami-gpu", wantCreated: "2026-08-01T00:00:00.000Z", wantMatches: 2,
},
{
name: "all images excluded",
images: []*ec2.Image{img("ami-gpu", "2026-08-01T00:00:00.000Z", "gpu-vendor")},
excludeKey: "gpu-vendor",
},
{
name: "nil tag entry does not panic",
images: []*ec2.Image{img("ami-a", "2026-01-01T00:00:00.000Z", "", "gpu-vendor"), img("ami-b", "2026-02-01T00:00:00.000Z", "")},
excludeKey: "gpu-vendor",
wantID: "ami-b", wantCreated: "2026-02-01T00:00:00.000Z", wantMatches: 1,
},
{
name: "missing creation date sorts oldest",
images: []*ec2.Image{&ec2.Image{ImageId: aws.String("ami-nodate")}, img("ami-dated", "2020-01-01T00:00:00.000Z")},
wantID: "ami-dated", wantCreated: "2020-01-01T00:00:00.000Z", wantMatches: 2,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
gotID, gotCreated, gotMatches := SelectNewestImage(tt.images, tt.excludeKey)
if gotID != tt.wantID {
t.Errorf("imageID = %q, want %q", gotID, tt.wantID)
}
if gotCreated != tt.wantCreated {
t.Errorf("created = %q, want %q", gotCreated, tt.wantCreated)
}
if gotMatches != tt.wantMatches {
t.Errorf("matches = %d, want %d", gotMatches, tt.wantMatches)
}
})
}
}
// rawFixture returns the bytes of the repo's known-good MBR fixture, so a compressed copy of
// it decodes to something validateDiskImagePath's file(1) sniff already accepts elsewhere.
func rawFixture(t *testing.T) []byte {