From bedf5fc11cea9cc6b92f37597fe8624d25b8d371 Mon Sep 17 00:00:00 2001
From: Tao Bao <tbao@google.com>
Date: Fri, 18 Nov 2016 12:01:26 -0800
Subject: updater: Refactor set_stage() and get_stage() functions.

Add read_bootloader_message_from() and write_bootloader_message_to() to
allow specifying the BCB device (/misc).

Also add testcases for set_stage() and get_stage().

Test: recovery_component_test passes.
Test: Build a recovery image and apply a two-step OTA package.
Change-Id: If5ab06a1aaaea168d2a9e5dd63c07c0a3190e4ae
---
 tests/Android.mk                 |  2 ++
 tests/component/updater_test.cpp | 59 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+)

(limited to 'tests')

diff --git a/tests/Android.mk b/tests/Android.mk
index fdc947028..5f6a7ce0c 100644
--- a/tests/Android.mk
+++ b/tests/Android.mk
@@ -80,10 +80,12 @@ LOCAL_STATIC_LIBRARIES := \
     libedify \
     libotafault \
     libupdater \
+    libbootloader_message \
     libverifier \
     libminui \
     libotautil \
     libmounts \
+    libfs_mgr \
     liblog \
     libselinux \
     libext4_utils_static \
diff --git a/tests/component/updater_test.cpp b/tests/component/updater_test.cpp
index cd285729d..f31f1f82a 100644
--- a/tests/component/updater_test.cpp
+++ b/tests/component/updater_test.cpp
@@ -23,6 +23,7 @@
 #include <android-base/file.h>
 #include <android-base/properties.h>
 #include <android-base/test_utils.h>
+#include <bootloader_message/bootloader_message.h>
 #include <gtest/gtest.h>
 #include <ziparchive/zip_archive.h>
 
@@ -451,3 +452,61 @@ TEST_F(UpdaterTest, write_value) {
   script = "write_value(\"value\", \"/proc/0/file1\")";
   expect("", script.c_str(), kNoCause);
 }
+
+TEST_F(UpdaterTest, get_stage) {
+  // get_stage() expects one argument.
+  expect(nullptr, "get_stage()", kArgsParsingFailure);
+  expect(nullptr, "get_stage(\"arg1\", \"arg2\")", kArgsParsingFailure);
+  expect(nullptr, "get_stage(\"arg1\", \"arg2\", \"arg3\")", kArgsParsingFailure);
+
+  // Set up a local file as BCB.
+  TemporaryFile tf;
+  std::string temp_file(tf.path);
+  bootloader_message boot;
+  strlcpy(boot.stage, "2/3", sizeof(boot.stage));
+  std::string err;
+  ASSERT_TRUE(write_bootloader_message_to(boot, temp_file, &err));
+
+  // Can read the stage value.
+  std::string script("get_stage(\"" + temp_file + "\")");
+  expect("2/3", script.c_str(), kNoCause);
+
+  // Bad BCB path.
+  script = "get_stage(\"doesntexist\")";
+  expect("", script.c_str(), kNoCause);
+}
+
+TEST_F(UpdaterTest, set_stage) {
+  // set_stage() expects two arguments.
+  expect(nullptr, "set_stage()", kArgsParsingFailure);
+  expect(nullptr, "set_stage(\"arg1\")", kArgsParsingFailure);
+  expect(nullptr, "set_stage(\"arg1\", \"arg2\", \"arg3\")", kArgsParsingFailure);
+
+  // Set up a local file as BCB.
+  TemporaryFile tf;
+  std::string temp_file(tf.path);
+  bootloader_message boot;
+  strlcpy(boot.command, "command", sizeof(boot.command));
+  strlcpy(boot.stage, "2/3", sizeof(boot.stage));
+  std::string err;
+  ASSERT_TRUE(write_bootloader_message_to(boot, temp_file, &err));
+
+  // Write with set_stage().
+  std::string script("set_stage(\"" + temp_file + "\", \"1/3\")");
+  expect(tf.path, script.c_str(), kNoCause);
+
+  // Verify.
+  bootloader_message boot_verify;
+  ASSERT_TRUE(read_bootloader_message_from(&boot_verify, temp_file, &err));
+
+  // Stage should be updated, with command part untouched.
+  ASSERT_STREQ("1/3", boot_verify.stage);
+  ASSERT_STREQ(boot.command, boot_verify.command);
+
+  // Bad BCB path.
+  script = "set_stage(\"doesntexist\", \"1/3\")";
+  expect("", script.c_str(), kNoCause);
+
+  script = "set_stage(\"/dev/full\", \"1/3\")";
+  expect("", script.c_str(), kNoCause);
+}
-- 
cgit v1.2.3