From 5c586101815ce71df2e55dddb8b4304cf50e4658 Mon Sep 17 00:00:00 2001
From: PJ <ipriyanshujain@gmail.com>
Date: Sun, 31 May 2026 15:27:45 +0530
Subject: [PATCH] test(runner): assert step uses Snapshot, not raw
 hierarchy/screenshot

TestRunner_UsesAtomicSnapshot catches regressions to the two-goroutine
race, and the existing parallel-fetch test now keys off ActionSnapshot.
---
 internal/runner/runner_test.go | 56 +++++++++++++++++++++++++++++++---
 1 file changed, 51 insertions(+), 5 deletions(-)

diff --git a/internal/runner/runner_test.go b/internal/runner/runner_test.go
index 4d433ab..79f3211 100644
--- a/internal/runner/runner_test.go
+++ b/internal/runner/runner_test.go
@@ -430,19 +430,19 @@ func TestRunner_ParallelFetchCallsAllDriverMethods(t *testing.T) {
 	}
 
 	actions := state.mock.Actions()
-	var hasHierarchy, hasMetrics, hasLogs bool
+	var hasSnapshot, hasMetrics, hasLogs bool
 	for _, a := range actions {
 		switch a.Kind {
-		case mockdriver.ActionHierarchy:
-			hasHierarchy = true
+		case mockdriver.ActionSnapshot:
+			hasSnapshot = true
 		case mockdriver.ActionMetrics:
 			hasMetrics = true
 		case mockdriver.ActionRecentLogs:
 			hasLogs = true
 		}
 	}
-	if !hasHierarchy {
-		t.Error("expected Hierarchy call in mock actions")
+	if !hasSnapshot {
+		t.Error("expected Snapshot call in mock actions")
 	}
 	if !hasMetrics {
 		t.Error("expected Metrics call in mock actions")
@@ -452,6 +452,52 @@ func TestRunner_ParallelFetchCallsAllDriverMethods(t *testing.T) {
 	}
 }
 
+// TestRunner_UsesAtomicSnapshot ensures the runner observes a step's UI
+// through the paired Snapshot RPC instead of racing two independent
+// hierarchy + screenshot reads. The pair must come from one on-device
+// frame; a regression to separate calls is what this test catches.
+func TestRunner_UsesAtomicSnapshot(t *testing.T) {
+	state := newHarness(t)
+	state.mock.ImageData = driver.Image{PNG: []byte("png"), Width: 1, Height: 1}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	summary, err := Run(ctx, Options{
+		Duration:    100 * time.Millisecond,
+		IdleTimeout: 20 * time.Millisecond,
+		Driver:      state.mock,
+		Verifier:    state.verifier,
+		TraceWriter: state.writer,
+	})
+	if err != nil {
+		t.Fatalf("Run: %v", err)
+	}
+	if summary.Steps == 0 {
+		t.Fatal("expected at least one step")
+	}
+
+	var snapshotCalls, hierarchyCalls, screenshotCalls int
+	for _, action := range state.mock.Actions() {
+		switch action.Kind {
+		case mockdriver.ActionSnapshot:
+			snapshotCalls++
+		case mockdriver.ActionHierarchy:
+			hierarchyCalls++
+		case mockdriver.ActionScreenshot:
+			screenshotCalls++
+		}
+	}
+	if snapshotCalls == 0 {
+		t.Errorf("expected at least one Snapshot call, got %d", snapshotCalls)
+	}
+	if hierarchyCalls != 0 {
+		t.Errorf("expected zero standalone Hierarchy calls (runner must use Snapshot), got %d", hierarchyCalls)
+	}
+	if screenshotCalls != 0 {
+		t.Errorf("expected zero standalone Screenshot calls (runner must use Snapshot), got %d", screenshotCalls)
+	}
+}
+
 // TestRunner_OneScreenshotPerStep verifies the runner writes a single
 // screenshot per step, captured concurrently with hierarchy so the two
 // observations describe the same UI moment.