diff --git a/internal/driver/ioscompanion/driver.go b/internal/driver/ioscompanion/driver.go index 2373df8..74a73cc 100644 --- a/internal/driver/ioscompanion/driver.go +++ b/internal/driver/ioscompanion/driver.go @@ -806,11 +806,15 @@ func (d *Driver) Snapshot(ctx context.Context) (string, driver.Image, error) { // hybrid path, so they are captured concurrently. Only the hierarchy leg // runs under withRecovery: two concurrent recoveries would race the // restart bookkeeping, and a screenshot connection failure surfaces as a - // plain error that the next serialized call recovers from. + // plain error that the next serialized call recovers from. The goroutine + // works through a captured local because a hierarchy-leg recovery + // reassigns d.companion mid-flight; a screenshot against the torn-down + // transport then fails as a plain error rather than racing the field. var data []byte screenshotDone := make(chan error, 1) + companion := d.companion go func() { - imageData, _, callErr := d.companion.Screenshot(ctx) + imageData, _, callErr := companion.Screenshot(ctx) data = imageData screenshotDone <- callErr }() diff --git a/internal/driver/ioscompanion/driver_test.go b/internal/driver/ioscompanion/driver_test.go index c1b290a..32ff96e 100644 --- a/internal/driver/ioscompanion/driver_test.go +++ b/internal/driver/ioscompanion/driver_test.go @@ -255,6 +255,49 @@ func TestSnapshotPairsHierarchyAndScreenshot(t *testing.T) { } } +// blockingScreenshotCompanion holds its Screenshot until proceed closes, so a +// test can keep the screenshot leg in flight while the hierarchy leg recovers. +type blockingScreenshotCompanion struct { + fakeCompanion + proceed chan struct{} +} + +func (b *blockingScreenshotCompanion) Screenshot(ctx context.Context) ([]byte, string, error) { + <-b.proceed + return b.fakeCompanion.Screenshot(ctx) +} + +func TestSnapshotRestartDuringScreenshotDoesNotRace(t *testing.T) { + // The hierarchy leg drops its connection, forcing withRecovery to restart + // while the screenshot goroutine is still in flight. The restart reassigns + // d.companion the way respawnAndRedial does; the goroutine must keep + // working through the transport it captured rather than racing the field. + first := &blockingScreenshotCompanion{ + fakeCompanion: fakeCompanion{ + accessibilityErr: status.Error(codes.Unavailable, "companion gone"), + screenshotData: samplePNG(t, 390, 844), + }, + proceed: make(chan struct{}), + } + replacement := &fakeCompanion{ + accessibilityJSON: "[]", + screenshotData: samplePNG(t, 390, 844), + } + d := newTestDriver(first) + d.restart = func(context.Context) error { + d.companion = replacement + close(first.proceed) + return nil + } + _, image, err := d.Snapshot(context.Background()) + if err != nil { + t.Fatalf("Snapshot should recover: %v", err) + } + if image.Width != 390 || image.Height != 844 { + t.Fatalf("image dims = %dx%d, want 390x844", image.Width, image.Height) + } +} + func TestScreenshotRejectsNonPNG(t *testing.T) { d := newTestDriver(&fakeCompanion{screenshotData: []byte("not a png")}) if _, err := d.Screenshot(context.Background()); err == nil {