Added Mock function call tests

Checking flash partition is called correct number of times

Test: tested fastboot_test
Change-Id: I8a04fb0e30e794641f3e2fbe2ab666274be18b85
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();
+    }
+}