add use-snapshot-before-first-module option

Bug: 322412784
Test: run ABTD with the option enabled
Change-Id: I6caee277c9dec798eb0577899de7b9a4acebb5b8

1 files changed, 12 insertions, 0 deletions

diff --git a/src/com/android/tradefed/testtype/suite/ITestSuite.java b/src/com/android/tradefed/testtype/suite/ITestSuite.java
index 563aa297c..a275fb3ab 100644
--- a/src/com/android/tradefed/testtype/suite/ITestSuite.java
+++ b/src/com/android/tradefed/testtype/suite/ITestSuite.java
@@ -358,6 +358,13 @@ public abstract class ITestSuite
             description = "Feature flag to use snapshot/restore instead of powerwash.")
     private boolean mUseSnapshotForReset = false;
 
+    @Option(
+            name = "use-snapshot-before-first-module",
+            description =
+                    "Immediately restore the device after taking a snapshot. Ensures tests are"
+                            + " consistently run within a restored VM.")
+    private boolean mUseSnapshotBeforeFirstModule = false;
+
     @Option(name = "stage-remote-file", description = "Whether to allow staging of remote files.")
     private boolean mStageRemoteFile = true;
 
@@ -792,6 +799,11 @@ public abstract class ITestSuite
                             .getSuiteSnapshots()
                             .put(mDevice, mContext.getInvocationId());
                 }
+                if (mUseSnapshotBeforeFirstModule) {
+                    String snapshot =
+                            ((AdbTcpConnection) connection).getSuiteSnapshots().get(mDevice);
+                    ((AdbTcpConnection) connection).recoverVirtualDevice(mDevice, snapshot, null);
+                }
             }
         }