2 files changed, 47 insertions, 3 deletions

diff --git a/tests/unit/bootloader_message_test.cpp b/tests/unit/bootloader_message_test.cpp
index b005d199c..95d875e69 100644
--- a/tests/unit/bootloader_message_test.cpp
+++ b/tests/unit/bootloader_message_test.cpp
@@ -15,6 +15,7 @@
  */
 
 #include <string>
+#include <string_view>
 #include <vector>
 
 #include <android-base/file.h>
@@ -22,6 +23,10 @@
 #include <bootloader_message/bootloader_message.h>
 #include <gtest/gtest.h>
 
+using namespace std::string_literals;
+
+extern void SetMiscBlockDeviceForTest(std::string_view misc_device);
+
 TEST(BootloaderMessageTest, read_and_write_bootloader_message) {
   TemporaryFile temp_misc;
 
@@ -114,3 +119,36 @@ TEST(BootloaderMessageTest, update_bootloader_message_recovery_options_long) {
             std::string(boot.reserved, sizeof(boot.reserved)));
 }
 
+TEST(BootloaderMessageTest, WriteMiscPartitionVendorSpace) {
+  TemporaryFile temp_misc;
+  ASSERT_TRUE(android::base::WriteStringToFile(std::string(4096, '\x00'), temp_misc.path));
+  SetMiscBlockDeviceForTest(temp_misc.path);
+
+  constexpr std::string_view kTestMessage = "kTestMessage";
+  std::string err;
+  ASSERT_TRUE(WriteMiscPartitionVendorSpace(kTestMessage.data(), kTestMessage.size(), 0, &err));
+
+  std::string message;
+  message.resize(kTestMessage.size());
+  ASSERT_TRUE(ReadMiscPartitionVendorSpace(message.data(), message.size(), 0, &err));
+  ASSERT_EQ(kTestMessage, message);
+
+  // Write with an offset.
+  ASSERT_TRUE(WriteMiscPartitionVendorSpace("\x00\x00", 2, 5, &err));
+  ASSERT_TRUE(ReadMiscPartitionVendorSpace(message.data(), message.size(), 0, &err));
+  ASSERT_EQ("kTest\x00\x00ssage"s, message);
+
+  // Write with the right size.
+  auto start_offset =
+      WIPE_PACKAGE_OFFSET_IN_MISC - VENDOR_SPACE_OFFSET_IN_MISC - kTestMessage.size();
+  ASSERT_TRUE(
+      WriteMiscPartitionVendorSpace(kTestMessage.data(), kTestMessage.size(), start_offset, &err));
+
+  // Out-of-bound write.
+  ASSERT_FALSE(WriteMiscPartitionVendorSpace(kTestMessage.data(), kTestMessage.size(),
+                                             start_offset + 1, &err));
+
+  // Message won't fit.
+  std::string long_message(WIPE_PACKAGE_OFFSET_IN_MISC - VENDOR_SPACE_OFFSET_IN_MISC + 1, 'a');
+  ASSERT_FALSE(WriteMiscPartitionVendorSpace(long_message.data(), long_message.size(), 0, &err));
+}
diff --git a/tests/unit/updater_test.cpp b/tests/unit/updater_test.cpp
index 4a8d1e6ff..81229f50a 100644
--- a/tests/unit/updater_test.cpp
+++ b/tests/unit/updater_test.cpp
@@ -52,13 +52,14 @@
 #include "updater/blockimg.h"
 #include "updater/install.h"
 #include "updater/updater.h"
+#include "updater/updater_runtime.h"
 
 using namespace std::string_literals;
 
 using PackageEntries = std::unordered_map<std::string, std::string>;
 
 static void expect(const char* expected, const std::string& expr_str, CauseCode cause_code,
-                   Updater* updater = nullptr) {
+                   Updater* updater) {
   std::unique_ptr<Expr> e;
   int error_count = 0;
   ASSERT_EQ(0, ParseString(expr_str, &e, &error_count));
@@ -83,6 +84,11 @@ static void expect(const char* expected, const std::string& expr_str, CauseCode
   ASSERT_EQ(cause_code, state.cause_code);
 }
 
+static void expect(const char* expected, const std::string& expr_str, CauseCode cause_code) {
+  Updater updater;
+  expect(expected, expr_str, cause_code, &updater);
+}
+
 static void BuildUpdatePackage(const PackageEntries& entries, int fd) {
   FILE* zip_file_ptr = fdopen(fd, "wb");
   ZipWriter zip_writer(zip_file_ptr);
@@ -168,9 +174,9 @@ class UpdaterTestBase {
 
     // Set up the handler, command_pipe, patch offset & length.
     TemporaryFile temp_pipe;
-    ASSERT_TRUE(updater_.Init(temp_pipe.release(), zip_file.path, false, nullptr));
+    ASSERT_TRUE(updater_.Init(temp_pipe.release(), zip_file.path, false));
     ASSERT_TRUE(updater_.RunUpdate());
-    ASSERT_EQ(result, updater_.result());
+    ASSERT_EQ(result, updater_.GetResult());
 
     // Parse the cause code written to the command pipe.
     int received_cause_code = kNoCause;