don't require "use_cvd" for CF snapshot support

Bug: 331861133
Test: TODO
Change-Id: I016cc0c8e71d7fa84c58d443db95dfceca24de7e

2 files changed, 49 insertions, 24 deletions

diff --git a/javatests/com/android/tradefed/device/connection/AdbSshConnectionTest.java b/javatests/com/android/tradefed/device/connection/AdbSshConnectionTest.java
index 87fbf8844..de7ad05fb 100644
--- a/javatests/com/android/tradefed/device/connection/AdbSshConnectionTest.java
+++ b/javatests/com/android/tradefed/device/connection/AdbSshConnectionTest.java
@@ -550,7 +550,6 @@ public class AdbSshConnectionTest {
         String snapshotId = "snapshot_user1";
         OptionSetter setter = new OptionSetter(mOptions);
         setter.setOptionValue("instance-user", instanceUser);
-        String cvdBin = String.format("/home/%s/bin/cvd", instanceUser);
         String snapshotPath = String.format("/tmp/%s/snapshots/%s", instanceUser, snapshotId);
         String snapshotCommandPath = String.format("--snapshot_path=%s", snapshotPath);
         String restoreSnapshotCommandPath = String.format("--snapshot_path=%s", snapshotPath);
@@ -609,8 +608,8 @@ public class AdbSshConnectionTest {
                         Mockito.eq("-i"),
                         Mockito.any(),
                         Mockito.eq(avdConnectHost),
-                        Mockito.eq(cvdBin),
-                        Mockito.eq("snapshot_take"),
+                        Mockito.eq(String.format("/home/%s/bin/snapshot_util_cvd", instanceUser)),
+                        Mockito.eq("--subcmd=snapshot_take"),
                         Mockito.eq("--force"),
                         Mockito.eq("--auto_suspend"),
                         Mockito.eq(snapshotCommandPath)))
@@ -631,8 +630,7 @@ public class AdbSshConnectionTest {
                         Mockito.eq("-i"),
                         Mockito.any(),
                         Mockito.eq(avdConnectHost),
-                        Mockito.eq(cvdBin),
-                        Mockito.eq("stop")))
+                        Mockito.eq(String.format("/home/%s/bin/stop_cvd", instanceUser))))
                 .thenReturn(successCmdResult);
         when(mMockRunUtil.runTimedCmd(
                         Mockito.anyLong(),
@@ -650,8 +648,7 @@ public class AdbSshConnectionTest {
                         Mockito.eq("-i"),
                         Mockito.any(),
                         Mockito.eq(avdConnectHost),
-                        Mockito.eq(cvdBin),
-                        Mockito.eq("start"),
+                        Mockito.eq(String.format("/home/%s/bin/launch_cvd", instanceUser)),
                         Mockito.eq(restoreSnapshotCommandPath)))
                 .thenReturn(successCmdResult);
         CommandResult adbResult = new CommandResult();
diff --git a/src/com/android/tradefed/device/connection/AdbSshConnection.java b/src/com/android/tradefed/device/connection/AdbSshConnection.java
index 9a4439c6a..8977cc643 100644
--- a/src/com/android/tradefed/device/connection/AdbSshConnection.java
+++ b/src/com/android/tradefed/device/connection/AdbSshConnection.java
@@ -667,15 +667,28 @@ public class AdbSshConnection extends AdbTcpConnection {
             user = getDevice().getOptions().getInstanceUser();
         }
 
-        String snapshotCommand =
-                commandBuilder(
-                        "cvd",
-                        String.format(
-                                "snapshot_take --force --auto_suspend"
-                                        + " --snapshot_path=/tmp/%s/snapshots/%s",
-                                user, snapshotId),
-                        user,
-                        offset);
+        String snapshotCommand;
+        if (getDevice().getOptions().getExtraOxygenArgs().containsKey("use_cvd")) {
+            snapshotCommand =
+                    commandBuilder(
+                            "cvd",
+                            String.format(
+                                    "snapshot_take --force --auto_suspend"
+                                            + " --snapshot_path=/tmp/%s/snapshots/%s",
+                                    user, snapshotId),
+                            user,
+                            offset);
+        } else {
+            snapshotCommand =
+                    commandBuilder(
+                            "snapshot_util_cvd",
+                            String.format(
+                                    "--subcmd=snapshot_take --force --auto_suspend"
+                                            + " --snapshot_path=/tmp/%s/snapshots/%s",
+                                    user, snapshotId),
+                            user,
+                            offset);
+        }
         if (Strings.isNullOrEmpty(snapshotCommand)) {
             throw new TargetSetupError(
                     "failed to set up snapshot command, invalid path",
@@ -733,13 +746,23 @@ public class AdbSshConnection extends AdbTcpConnection {
             user = getDevice().getOptions().getInstanceUser();
         }
 
-        String restoreCommand =
-                commandBuilder(
-                        "cvd",
-                        String.format(
-                                "start --snapshot_path=/tmp/%s/snapshots/%s", user, snapshotId),
-                        user,
-                        offset);
+        String restoreCommand;
+        if (getDevice().getOptions().getExtraOxygenArgs().containsKey("use_cvd")) {
+            restoreCommand =
+                    commandBuilder(
+                            "cvd",
+                            String.format(
+                                    "start --snapshot_path=/tmp/%s/snapshots/%s", user, snapshotId),
+                            user,
+                            offset);
+        } else {
+            restoreCommand =
+                    commandBuilder(
+                            "launch_cvd",
+                            String.format("--snapshot_path=/tmp/%s/snapshots/%s", user, snapshotId),
+                            user,
+                            offset);
+        }
         if (restoreCommand.length() == 0) {
             throw new TargetSetupError(
                     "failed to set up restore command, invalid path",
@@ -804,7 +827,12 @@ public class AdbSshConnection extends AdbTcpConnection {
             user = getDevice().getOptions().getInstanceUser();
         }
 
-        String stopCommand = commandBuilder("cvd", "stop", user, offset);
+        String stopCommand;
+        if (getDevice().getOptions().getExtraOxygenArgs().containsKey("use_cvd")) {
+            stopCommand = commandBuilder("cvd", "stop", user, offset);
+        } else {
+            stopCommand = commandBuilder("stop_cvd", "", user, offset);
+        }
         if (stopCommand.length() == 0) {
             throw new TargetSetupError(
                     "failed to set up stop command, invalid path",