1 files changed, 30 insertions, 19 deletions

diff --git a/update_manager/update_manager.h b/update_manager/update_manager.h
index b0fd97fa..e266b574 100644
--- a/update_manager/update_manager.h
+++ b/update_manager/update_manager.h
@@ -22,10 +22,9 @@
 #include <string>
 
 #include <base/callback.h>
-#include <base/memory/ref_counted.h>
 #include <base/time/time.h>
 
-#include "update_engine/common/clock_interface.h"
+#include "update_engine/common/system_state.h"
 #include "update_engine/update_manager/default_policy.h"
 #include "update_engine/update_manager/evaluation_context.h"
 #include "update_engine/update_manager/policy.h"
@@ -33,22 +32,31 @@
 
 namespace chromeos_update_manager {
 
-// Comparator for scoped_refptr objects.
-template <typename T>
-struct ScopedRefPtrLess {
-  bool operator()(const scoped_refptr<T>& first,
-                  const scoped_refptr<T>& second) const {
-    return first.get() < second.get();
-  }
+// Please do not move this class into a new file for simplicity.
+// This pure virtual class is purely created for purpose of testing. The reason
+// was that |UpdateManager|'s member functions are templatized, which does not
+// play nicely when testing (mocking + faking). Whenever a specialized member of
+// |UpdateManager| must be tested, please add a specialized template member
+// function within this class for testing.
+class SpecializedPolicyRequestInterface {
+ public:
+  virtual ~SpecializedPolicyRequestInterface() = default;
+
+  virtual void AsyncPolicyRequestUpdateCheckAllowed(
+      base::Callback<void(EvalStatus, const UpdateCheckParams& result)>
+          callback,
+      EvalStatus (Policy::*policy_method)(EvaluationContext*,
+                                          State*,
+                                          std::string*,
+                                          UpdateCheckParams*) const) = 0;
 };
 
 // The main Update Manager singleton class.
-class UpdateManager {
+class UpdateManager : public SpecializedPolicyRequestInterface {
  public:
   // Creates the UpdateManager instance, assuming ownership on the provided
   // |state|.
-  UpdateManager(chromeos_update_engine::ClockInterface* clock,
-                base::TimeDelta evaluation_timeout,
+  UpdateManager(base::TimeDelta evaluation_timeout,
                 base::TimeDelta expiration_timeout,
                 State* state);
 
@@ -91,6 +99,14 @@ class UpdateManager {
           EvaluationContext*, State*, std::string*, R*, ExpectedArgs...) const,
       ActualArgs... args);
 
+  void AsyncPolicyRequestUpdateCheckAllowed(
+      base::Callback<void(EvalStatus, const UpdateCheckParams& result)>
+          callback,
+      EvalStatus (Policy::*policy_method)(EvaluationContext*,
+                                          State*,
+                                          std::string*,
+                                          UpdateCheckParams*) const) override;
+
  protected:
   // The UpdateManager receives ownership of the passed Policy instance.
   void set_policy(const Policy* policy) { policy_.reset(policy); }
@@ -125,7 +141,7 @@ class UpdateManager {
   // the evaluation will be re-scheduled to be called later.
   template <typename R, typename... Args>
   void OnPolicyReadyToEvaluate(
-      scoped_refptr<EvaluationContext> ec,
+      std::shared_ptr<EvaluationContext> ec,
       base::Callback<void(EvalStatus status, const R& result)> callback,
       EvalStatus (Policy::*policy_method)(
           EvaluationContext*, State*, std::string*, R*, Args...) const,
@@ -145,9 +161,6 @@ class UpdateManager {
   // State Providers.
   std::unique_ptr<State> state_;
 
-  // Pointer to the mockable clock interface;
-  chromeos_update_engine::ClockInterface* clock_;
-
   // Timeout for a policy evaluation.
   const base::TimeDelta evaluation_timeout_;
 
@@ -159,9 +172,7 @@ class UpdateManager {
   // destructed; alternatively, when the UpdateManager instance is destroyed, it
   // will remove all pending events associated with all outstanding contexts
   // (which should, in turn, trigger their destruction).
-  std::set<scoped_refptr<EvaluationContext>,
-           ScopedRefPtrLess<EvaluationContext>>
-      ec_repo_;
+  std::set<std::shared_ptr<EvaluationContext>> ec_repo_;
 
   base::WeakPtrFactory<UpdateManager> weak_ptr_factory_;