1 files changed, 6 insertions, 8 deletions
diff --git a/update_boot_flags_action_unittest.cc b/update_boot_flags_action_unittest.cc
index 1b2bfa5a..26cbe909 100644
--- a/update_boot_flags_action_unittest.cc
+++ b/update_boot_flags_action_unittest.cc
@@ -22,18 +22,17 @@
 #include <base/bind.h>
 #include <gtest/gtest.h>
 
-#include "update_engine/fake_system_state.h"
+#include "update_engine/common/fake_boot_control.h"
 
 namespace chromeos_update_engine {
 
 class UpdateBootFlagsActionTest : public ::testing::Test {
- public:
-  FakeSystemState fake_system_state_;
+ protected:
+  FakeBootControl boot_control_;
 };
 
 TEST_F(UpdateBootFlagsActionTest, SimpleTest) {
-  auto boot_control = fake_system_state_.fake_boot_control();
-  auto action = std::make_unique<UpdateBootFlagsAction>(boot_control);
+  auto action = std::make_unique<UpdateBootFlagsAction>(&boot_control_);
   ActionProcessor processor;
   processor.EnqueueAction(std::move(action));
 
@@ -49,9 +48,8 @@ TEST_F(UpdateBootFlagsActionTest, DoubleActionTest) {
   UpdateBootFlagsAction::updated_boot_flags_ = false;
   UpdateBootFlagsAction::is_running_ = false;
 
-  auto boot_control = fake_system_state_.fake_boot_control();
-  auto action1 = std::make_unique<UpdateBootFlagsAction>(boot_control);
-  auto action2 = std::make_unique<UpdateBootFlagsAction>(boot_control);
+  auto action1 = std::make_unique<UpdateBootFlagsAction>(&boot_control_);
+  auto action2 = std::make_unique<UpdateBootFlagsAction>(&boot_control_);
   ActionProcessor processor1, processor2;
   processor1.EnqueueAction(std::move(action1));
   processor2.EnqueueAction(std::move(action2));