From 461cdc559374c2dd8234f1b50f8236f7caffc761 Mon Sep 17 00:00:00 2001
From: Julian Sommer <115516918+juliansommer@users.noreply.github.com>
Date: Mon, 24 Aug 2026 10:11:32 +1000
Subject: [PATCH] fix: improve tests

---
 .../ec2/instance/service_impl_test.go         | 54 ++++++++++++-------
 spinifex/vm/crash_recovery_test.go            | 24 +++++++--
 2 files changed, 57 insertions(+), 21 deletions(-)

diff --git a/spinifex/handlers/ec2/instance/service_impl_test.go b/spinifex/handlers/ec2/instance/service_impl_test.go
index df56e66de..06d70a6e0 100644
--- a/spinifex/handlers/ec2/instance/service_impl_test.go
+++ b/spinifex/handlers/ec2/instance/service_impl_test.go
@@ -435,34 +435,52 @@ func TestRunInstance_WithTags(t *testing.T) {
 					{Key: aws.String("env"), Value: aws.String("dev")},
 				},
 			},
+			{
+				ResourceType: aws.String("volume"),
+				Tags:         []*ec2.Tag{{Key: aws.String("scope"), Value: aws.String("volume-only")}},
+			},
 		},
 	}
 
-	instance, _, err := svc.RunInstance(input)
+	instance, ec2Instance, err := svc.RunInstance(input)
 	require.NoError(t, err)
-	// Tags are stored in RunInstancesInput which is preserved on the VM
-	assert.Equal(t, input, instance.RunInstancesInput)
-	assert.Len(t, input.TagSpecifications, 1)
-	assert.Len(t, input.TagSpecifications[0].Tags, 2)
+
+	// DescribeInstances and tag filters read ec2Instance.Tags, not the echoed
+	// input, and only instance-scoped specs may land there.
+	got := map[string]string{}
+	for _, tag := range ec2Instance.Tags {
+		got[aws.StringValue(tag.Key)] = aws.StringValue(tag.Value)
+	}
+	assert.Equal(t, map[string]string{"Name": "test-vm", "env": "dev"}, got)
+	assert.Equal(t, ec2Instance.Tags, instance.Instance.Tags)
 }
 
-func TestRunInstance_WithPlacement(t *testing.T) {
-	instanceTypes := map[string]*ec2.InstanceTypeInfo{
-		"t3.micro": {InstanceType: aws.String("t3.micro")},
-	}
-	svc := &InstanceServiceImpl{instanceTypes: instanceTypes}
-
-	input := &ec2.RunInstancesInput{
-		ImageId:      aws.String("ami-012345"),
-		InstanceType: aws.String("t3.micro"),
-		Placement: &ec2.Placement{
-			GroupName: aws.String("my-pg"),
+// Placement is projected onto the VM by PrepareRunInstances, not RunInstance,
+// so the assertion has to be driven through the entry point that reads it.
+func TestPrepareRunInstances_ProjectsPlacementGroup(t *testing.T) {
+	instanceTypes, _ := defaultPrepareInstanceTypes()
+	svc := &InstanceServiceImpl{
+		config:        &config.Config{},
+		instanceTypes: instanceTypes,
+		amiLoader: &fakeAMILoader{byID: map[string]ebsmetadata.AMI{
+			"ami-1": {ImageOwnerAlias: "acc", PlatformDetails: "Linux/UNIX"},
+		}},
+		resourceMgr: &fakeResourceCapacityProvider{
+			instanceTypes: instanceTypes,
+			canAllocFn:    func(_ *ec2.InstanceTypeInfo, count int) int { return count },
 		},
 	}
 
-	instance, _, err := svc.RunInstance(input)
+	_, instances, _, err := svc.PrepareRunInstances(context.Background(), &ec2.RunInstancesInput{
+		InstanceType: aws.String("t3.micro"),
+		ImageId:      aws.String("ami-1"),
+		MinCount:     aws.Int64(1),
+		MaxCount:     aws.Int64(1),
+		Placement:    &ec2.Placement{GroupName: aws.String("my-pg")},
+	}, "acc", "")
 	require.NoError(t, err)
-	assert.Equal(t, "my-pg", *instance.RunInstancesInput.Placement.GroupName)
+	require.Len(t, instances, 1)
+	assert.Equal(t, "my-pg", instances[0].PlacementGroupName)
 }
 
 func TestParseVolumeParams_MultipleBlockDeviceMappings(t *testing.T) {
diff --git a/spinifex/vm/crash_recovery_test.go b/spinifex/vm/crash_recovery_test.go
index 69e9322a1..226d29e5a 100644
--- a/spinifex/vm/crash_recovery_test.go
+++ b/spinifex/vm/crash_recovery_test.go
@@ -380,8 +380,24 @@ func TestMaybeRestart_ResetsAfterWindow(t *testing.T) {
 		"restart count must reset to 0 once the window expires")
 }
 
+// countScheduledRestarts replaces the restart seam with a counter. RestartCount
+// only moves inside the scheduled callback, so a guard that refuses to restart
+// is observable here and nowhere on the instance itself.
+func countScheduledRestarts(t *testing.T) *atomic.Int32 {
+	t.Helper()
+	var scheduled atomic.Int32
+	prev := restartAfterFunc
+	t.Cleanup(func() { restartAfterFunc = prev })
+	restartAfterFunc = func(_ time.Duration, _ func()) *time.Timer {
+		scheduled.Add(1)
+		return time.NewTimer(time.Hour)
+	}
+	return &scheduled
+}
+
 func TestMaybeRestart_SkipsShuttingDown(t *testing.T) {
 	m, _, _, shuttingDown := crashTestManager(t)
+	scheduled := countScheduledRestarts(t)
 	shuttingDown.Store(true)
 
 	instance := &VM{
@@ -397,11 +413,12 @@ func TestMaybeRestart_SkipsShuttingDown(t *testing.T) {
 
 	m.MaybeRestart(instance)
 
-	assert.Equal(t, 0, instance.Health.RestartCount, "shutdown skip must not increment restart")
+	assert.Zero(t, scheduled.Load(), "shutdown skip must not schedule a restart")
 }
 
 func TestMaybeRestart_UnknownInstanceType(t *testing.T) {
 	m, _, _, _ := crashTestManager(t)
+	scheduled := countScheduledRestarts(t)
 
 	instance := &VM{
 		ID:           "i-restart-unknown",
@@ -416,12 +433,13 @@ func TestMaybeRestart_UnknownInstanceType(t *testing.T) {
 
 	m.MaybeRestart(instance)
 
-	assert.Equal(t, 0, instance.Health.RestartCount,
+	assert.Zero(t, scheduled.Load(),
 		"unknown instance type must not schedule a restart")
 }
 
 func TestMaybeRestart_InsufficientResources(t *testing.T) {
 	m, rc, _, _ := crashTestManager(t)
+	scheduled := countScheduledRestarts(t)
 	rc.canAllocateRet = 0
 
 	instance := &VM{
@@ -437,7 +455,7 @@ func TestMaybeRestart_InsufficientResources(t *testing.T) {
 
 	m.MaybeRestart(instance)
 
-	assert.Equal(t, 0, instance.Health.RestartCount,
+	assert.Zero(t, scheduled.Load(),
 		"insufficient resources must not schedule a restart")
 }