From c57be03b6978135122cfd1d15394fb232886e0b8 Mon Sep 17 00:00:00 2001
From: PJ <ipriyanshujain@gmail.com>
Date: Sun, 31 May 2026 15:26:03 +0530
Subject: [PATCH] feat(driver): add Snapshot to chrome and mock drivers

The chrome tab is single-threaded so its Snapshot pairs the two reads
without extra locking. The mock records ActionSnapshot so tests can
assert the runner reaches for the paired RPC.
---
 internal/driver/chrome/driver.go | 15 +++++++++++++++
 internal/driver/mock/mock.go     | 15 +++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/internal/driver/chrome/driver.go b/internal/driver/chrome/driver.go
index f35c4a8..cb59da4 100644
--- a/internal/driver/chrome/driver.go
+++ b/internal/driver/chrome/driver.go
@@ -268,6 +268,21 @@ func (d *Driver) Screenshot(_ context.Context) (driver.Image, error) {
 	return driver.Image{PNG: buf, Width: w, Height: h}, nil
 }
 
+// Snapshot pairs hierarchy and screenshot back-to-back. The chromedp tab
+// is single-threaded so the two CDP round-trips are already serialized:
+// pairing them here matches the DeviceDriver contract without extra locking.
+func (d *Driver) Snapshot(ctx context.Context) (string, driver.Image, error) {
+	hierarchy, err := d.Hierarchy(ctx)
+	if err != nil {
+		return "", driver.Image{}, err
+	}
+	image, err := d.Screenshot(ctx)
+	if err != nil {
+		return hierarchy, driver.Image{}, err
+	}
+	return hierarchy, image, nil
+}
+
 func (d *Driver) RecentLogs(_ context.Context, since time.Time, minLevel string) ([]driver.LogEntry, error) {
 	sinceMillis := since.UnixMilli()
 	d.logsMu.Lock()
diff --git a/internal/driver/mock/mock.go b/internal/driver/mock/mock.go
index bafca5a..d03af0a 100644
--- a/internal/driver/mock/mock.go
+++ b/internal/driver/mock/mock.go
@@ -20,6 +20,7 @@ const (
 	ActionPressKey    ActionKind = "press_key"
 	ActionHierarchy   ActionKind = "hierarchy"
 	ActionScreenshot  ActionKind = "screenshot"
+	ActionSnapshot    ActionKind = "snapshot"
 	ActionRecentLogs  ActionKind = "recent_logs"
 	ActionWaitForIdle ActionKind = "wait_for_idle"
 	ActionHealth      ActionKind = "health"
@@ -234,6 +235,20 @@ func (d *Driver) Screenshot(ctx context.Context) (driver.Image, error) {
 	return d.ImageData, nil
 }
 
+// Snapshot returns the hierarchy + screenshot pair atomically, mirroring
+// the real driver's contract. It records a single ActionSnapshot so tests
+// can assert the runner reached for the paired RPC instead of racing the
+// two reads.
+func (d *Driver) Snapshot(ctx context.Context) (string, driver.Image, error) {
+	if err := d.failure(ActionSnapshot); err != nil {
+		return "", driver.Image{}, err
+	}
+	d.record(Action{Kind: ActionSnapshot})
+	d.mutex.Lock()
+	defer d.mutex.Unlock()
+	return d.HierarchyJSON, d.ImageData, nil
+}
+
 func (d *Driver) WaitForIdle(ctx context.Context, duration time.Duration) error {
 	if err := d.failure(ActionWaitForIdle); err != nil {
 		return err