From 0d2d631d275b49b91aa7cc0bc2dd6755f1058b7b Mon Sep 17 00:00:00 2001
From: Daniel Zheng <zhengdaniel@google.com>
Date: Tue, 11 Apr 2023 09:09:35 -0700
Subject: [PATCH] Added Mock function call tests

Checking flash partition is called correct number of times

Test: tested fastboot_test
Change-Id: I8a04fb0e30e794641f3e2fbe2ab666274be18b85
---
 fastboot/fastboot_driver_mock.h |  3 +--
 fastboot/task_test.cpp          | 16 ++++++++++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/fastboot/fastboot_driver_mock.h b/fastboot/fastboot_driver_mock.h
index 62a570849..d2a123b19 100644
--- a/fastboot/fastboot_driver_mock.h
+++ b/fastboot/fastboot_driver_mock.h
@@ -22,8 +22,7 @@ namespace fastboot {
 
 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 125226c0d..6fc20562a 100644
--- a/fastboot/task_test.cpp
+++ b/fastboot/task_test.cpp
@@ -136,3 +136,19 @@ TEST_F(ParseTest, CorrectTaskFormed) {
     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();
+    }
+}