diff --git a/aosp/update_engine_client_android.cc b/aosp/update_engine_client_android.cc
index 8628463..c00e9f5 100644
--- a/aosp/update_engine_client_android.cc
+++ b/aosp/update_engine_client_android.cc
@@ -202,6 +202,19 @@
     return ExitWhenIdle(1);
   }
 
+  // Other commands, such as |setShouldSwitchSlotOnReboot|, might rely on the
+  // follow behavior, so created callback before running these commands.
+  if (FLAGS_follow) {
+    // Register a callback object with the service.
+    callback_ = new UECallback(this);
+    bool bound;
+    if (!service_->bind(callback_, &bound).isOk() || !bound) {
+      LOG(ERROR) << "Failed to bind() the UpdateEngine daemon.";
+      return 1;
+    }
+    keep_running = true;
+  }
+
   if (FLAGS_suspend) {
     return ExitWhenIdle(service_->suspend());
   }
@@ -229,10 +242,15 @@
     if (should_switch) {
       status = service_->setShouldSwitchSlotOnReboot(
           android::String16(FLAGS_metadata.c_str(), FLAGS_metadata.size()));
+      if (!FLAGS_follow) {
+        return ExitWhenIdle(status);
+      }
     } else {
+      // resetShouldSwitchSlotOnReboot() is a synchronous call, no need to
+      // follow
       status = service_->resetShouldSwitchSlotOnReboot();
+      return ExitWhenIdle(status);
     }
-    return ExitWhenIdle(status);
   }
 
   if (FLAGS_verify) {
@@ -274,17 +292,6 @@
     keep_running = true;
   }
 
-  if (FLAGS_follow) {
-    // Register a callback object with the service.
-    callback_ = new UECallback(this);
-    bool bound;
-    if (!service_->bind(callback_, &bound).isOk() || !bound) {
-      LOG(ERROR) << "Failed to bind() the UpdateEngine daemon.";
-      return 1;
-    }
-    keep_running = true;
-  }
-
   if (FLAGS_update) {
     auto and_headers = ParseHeaders(FLAGS_headers);
     Status status = service_->applyPayload(
diff --git a/scripts/update_device.py b/scripts/update_device.py
index 165bc97..db653dc 100755
--- a/scripts/update_device.py
+++ b/scripts/update_device.py
@@ -438,6 +438,10 @@
                       help='Verify metadata then exit, instead of applying the OTA.')
   parser.add_argument('--no-care-map', action='store_true',
                       help='Do not push care_map.pb to device.')
+  parser.add_argument('--perform-slot-switch', action='store_true',
+                      help='Perform slot switch for this OTA package')
+  parser.add_argument('--perform-reset-slot-switch', action='store_true',
+                      help='Perform reset slot switch for this OTA package')
   args = parser.parse_args()
   logging.basicConfig(
       level=logging.WARNING if args.no_verbose else logging.INFO)
@@ -474,6 +478,16 @@
     # Return 0, as we are executing ADB commands here, no work needed after
     # this point
     return 0
+  if args.perform_slot_switch:
+    assert PushMetadata(dut, args.otafile, metadata_path)
+    dut.adb(["shell", "update_engine_client",
+            "--switch_slot=true", "--metadata={}".format(metadata_path), "--follow"])
+    return 0
+  if args.perform_reset_slot_switch:
+    assert PushMetadata(dut, args.otafile, metadata_path)
+    dut.adb(["shell", "update_engine_client",
+            "--switch_slot=false", "--metadata={}".format(metadata_path)])
+    return 0
 
   if args.no_slot_switch:
     args.extra_headers += "\nSWITCH_SLOT_ON_REBOOT=0"