diff --git a/fastboot/fastboot_driver_mock.h b/fastboot/fastboot_driver_mock.h
index 62a5708..d2a123b 100644
--- a/fastboot/fastboot_driver_mock.h
+++ b/fastboot/fastboot_driver_mock.h
@@ -22,8 +22,7 @@
 
 class MockFastbootDriver : public IFastBootDriver {
   public:
-    MOCK_METHOD(RetCode, FlashPartition,
-                (const std::string& partition, android::base::borrowed_fd fd, uint32_t sz),
+    MOCK_METHOD(RetCode, FlashPartition, (const std::string&, android::base::borrowed_fd, uint32_t),
                 (override));
     MOCK_METHOD(RetCode, DeletePartition, (const std::string&), (override));
     MOCK_METHOD(RetCode, Reboot, (std::string*, std::vector<std::string>*), (override));
diff --git a/fastboot/task_test.cpp b/fastboot/task_test.cpp
index 125226c..6fc2056 100644
--- a/fastboot/task_test.cpp
+++ b/fastboot/task_test.cpp
@@ -136,3 +136,19 @@
     ASSERT_TRUE(tasks[3]->AsUpdateSuperTask());
     ASSERT_TRUE(tasks[4]->AsWipeTask());
 }
+
+TEST_F(ParseTest, CorrectDriverCalls) {
+    fastboot::MockFastbootDriver fb;
+    fp->fb = &fb;
+
+    EXPECT_CALL(fb, RebootTo(_, _, _)).Times(1);
+    EXPECT_CALL(fb, Reboot(_, _)).Times(1);
+    EXPECT_CALL(fb, WaitForDisconnect()).Times(2);
+
+    std::vector<std::string> commands = {"reboot bootloader", "reboot"};
+    std::vector<std::unique_ptr<Task>> tasks = collectTasks(fp.get(), commands);
+
+    for (auto& task : tasks) {
+        task->Run();
+    }
+}