arc: Add wake lock service

This change adds wake lock service to acquire and release cpu wake
locks.

BUG=68944053
TEST=Verified cpu wake lock acquisition and release.

Change-Id: I7353ee394e9a400b4a32d182e551700162ce15a5
Reviewed-on: https://chromium-review.googlesource.com/936051Reviewed-by: John Abd-El-Malek <jam@chromium.org>
Reviewed-by: Luis Hector Chavez <lhchavez@chromium.org>
Reviewed-by: Jorge Lucangeli Obes <jorgelo@chromium.org>
Reviewed-by: Dan Erat <derat@chromium.org>
Commit-Queue: Abhishek Bhardwaj <abhishekbh@google.com>
Cr-Commit-Position: refs/heads/master@{#542683}

chrome/browser/chromeos/arc/arc_service_launcher.cc

@@ -59,6 +59,7 @@
 #include "components/arc/timer/arc_timer_bridge.h"
 #include "components/arc/usb/usb_host_bridge.h"
 #include "components/arc/volume_mounter/arc_volume_mounter_bridge.h"
+#include "components/arc/wake_lock/arc_wake_lock_bridge.h"
 #include "components/prefs/pref_member.h"
 #include "ui/arc/notification/arc_notification_manager.h"
 
@@ -176,6 +177,7 @@ void ArcServiceLauncher::OnPrimaryUserProfilePrepared(Profile* profile) {
   ArcVoiceInteractionArcHomeService::GetForBrowserContext(profile);
   ArcVoiceInteractionFrameworkService::GetForBrowserContext(profile);
   ArcVolumeMounterBridge::GetForBrowserContext(profile);
+  ArcWakeLockBridge::GetForBrowserContext(profile);
   ArcWallpaperService::GetForBrowserContext(profile);
   GpuArcVideoServiceHost::GetForBrowserContext(profile);
 

components/arc/BUILD.gn

@@ -61,6 +61,8 @@ static_library("arc") {
     "usb/usb_host_ui_delegate.h",
     "volume_mounter/arc_volume_mounter_bridge.cc",
     "volume_mounter/arc_volume_mounter_bridge.h",
+    "wake_lock/arc_wake_lock_bridge.cc",
+    "wake_lock/arc_wake_lock_bridge.h",
   ]
 
   public_deps = [
@@ -202,6 +204,8 @@ static_library("arc_test_support") {
     "test/fake_timer_instance.h",
     "test/fake_voice_interaction_framework_instance.cc",
     "test/fake_voice_interaction_framework_instance.h",
+    "test/fake_wake_lock_instance.cc",
+    "test/fake_wake_lock_instance.h",
     "test/fake_wallpaper_instance.cc",
     "test/fake_wallpaper_instance.h",
     "test/test_browser_context.cc",
@@ -242,6 +246,7 @@ source_set("unit_tests") {
     "metrics/arc_metrics_service_unittest.cc",
     "power/arc_power_bridge_unittest.cc",
     "timer/arc_timer_bridge_unittest.cc",
+    "wake_lock/arc_wake_lock_bridge_unittest.cc",
   ]
 
   deps = [

components/arc/arc_bridge_host_impl.cc

@@ -287,6 +287,11 @@ void ArcBridgeHostImpl::OnVolumeMounterInstanceReady(
                   std::move(volume_mounter_ptr));
 }
 
+void ArcBridgeHostImpl::OnWakeLockInstanceReady(
+    mojom::WakeLockInstancePtr wakelock_ptr) {
+  OnInstanceReady(arc_bridge_service_->wake_lock(), std::move(wakelock_ptr));
+}
+
 void ArcBridgeHostImpl::OnWallpaperInstanceReady(
     mojom::WallpaperInstancePtr wallpaper_ptr) {
   OnInstanceReady(arc_bridge_service_->wallpaper(), std::move(wallpaper_ptr));

components/arc/arc_bridge_host_impl.h

@@ -97,6 +97,8 @@ class ArcBridgeHostImpl : public mojom::ArcBridgeHost {
       mojom::VoiceInteractionFrameworkInstancePtr framework_ptr) override;
   void OnVolumeMounterInstanceReady(
       mojom::VolumeMounterInstancePtr volume_mounter_ptr) override;
+  void OnWakeLockInstanceReady(
+      mojom::WakeLockInstancePtr wake_lock_ptr) override;
   void OnWallpaperInstanceReady(
       mojom::WallpaperInstancePtr wallpaper_ptr) override;
 

components/arc/arc_bridge_service.h

@@ -83,6 +83,8 @@ class VoiceInteractionFrameworkHost;
 class VoiceInteractionFrameworkInstance;
 class VolumeMounterHost;
 class VolumeMounterInstance;
+class WakeLockHost;
+class WakeLockInstance;
 class WallpaperHost;
 class WallpaperInstance;
 
@@ -218,6 +220,9 @@ class ArcBridgeService {
   volume_mounter() {
     return &volume_mounter_;
   }
+  ConnectionHolder<mojom::WakeLockInstance, mojom::WakeLockHost>* wake_lock() {
+    return &wake_lock_;
+  }
   ConnectionHolder<mojom::WallpaperInstance, mojom::WallpaperHost>*
   wallpaper() {
     return &wallpaper_;
@@ -278,6 +283,7 @@ class ArcBridgeService {
       voice_interaction_framework_;
   ConnectionHolder<mojom::VolumeMounterInstance, mojom::VolumeMounterHost>
       volume_mounter_;
+  ConnectionHolder<mojom::WakeLockInstance, mojom::WakeLockHost> wake_lock_;
   ConnectionHolder<mojom::WallpaperInstance, mojom::WallpaperHost> wallpaper_;
 
   DISALLOW_COPY_AND_ASSIGN(ArcBridgeService);

components/arc/common/BUILD.gn

@@ -52,6 +52,7 @@ if (is_chromeos) {
       "voice_interaction_arc_home.mojom",
       "voice_interaction_framework.mojom",
       "volume_mounter.mojom",
+      "wake_lock.mojom",
       "wallpaper.mojom",
     ]
 

components/arc/common/arc_bridge.mojom

@@ -42,11 +42,12 @@ import "components/arc/common/video.mojom";
 import "components/arc/common/voice_interaction_arc_home.mojom";
 import "components/arc/common/voice_interaction_framework.mojom";
 import "components/arc/common/volume_mounter.mojom";
+import "components/arc/common/wake_lock.mojom";
 import "components/arc/common/wallpaper.mojom";
 
-// Next MinVersion: 37
+// Next MinVersion: 38
 // Deprecated method IDs: 101, 105
-// Next method ID: 142
+// Next method ID: 143
 interface ArcBridgeHost {
   // Keep the entries alphabetical. In order to do so without breaking
   // compatibility with the ARC instance, explicitly assign each interface a
@@ -176,6 +177,9 @@ interface ArcBridgeHost {
   // Notifies Chrome that the VolumeMounter interface is ready.
   [MinVersion=25] OnVolumeMounterInstanceReady@131(VolumeMounterInstance instance_ptr);
 
+  // Notifies Chrome that the WakeLockInstance interface is ready.
+  [MinVersion=37] OnWakeLockInstanceReady@142(WakeLockInstance instance_ptr);
+
   // Notifies Chrome that the WallpaperInstance interface is ready.
   [MinVersion=18] OnWallpaperInstanceReady@124(WallpaperInstance instance_ptr);
 };

components/arc/common/wake_lock.mojom

+// Copyright 2018 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+// Next MinVersion: 1
+
+module arc.mojom;
+
+// Interface for the mojo service to acquire and release wake locks on behalf of
+// the instance. On a connection error, all wake locks are released.
+//
+// Next method ID: 2
+interface WakeLockHost {
+  // Increases the reference count for the CPU wake lock on the host. Acquires
+  // the CPU wake lock when reference count is incremented to one. Returns true
+  // in all cases.
+  AcquirePartialWakeLock@0() => (bool result);
+
+  // Decreases the reference count for the CPU wake lock on the host. Releases
+  // the wake lock when reference count is decremented to zero. Does nothing
+  // if reference count is zero at the time of call. Returns true in all cases.
+  ReleasePartialWakeLock@1() => (bool result);
+};
+
+// Next method ID: 1
+interface WakeLockInstance {
+  // Establishes full-duplex communication with the host.
+  Init@0(WakeLockHost host_ptr) => ();
+};

components/arc/test/fake_arc_bridge_host.cc

@@ -119,6 +119,9 @@ void FakeArcBridgeHost::OnVoiceInteractionFrameworkInstanceReady(
 void FakeArcBridgeHost::OnVolumeMounterInstanceReady(
     mojom::VolumeMounterInstancePtr volume_mounter_ptr) {}
 
+void FakeArcBridgeHost::OnWakeLockInstanceReady(
+    mojom::WakeLockInstancePtr wakelock_ptr) {}
+
 void FakeArcBridgeHost::OnWallpaperInstanceReady(
     mojom::WallpaperInstancePtr wallpaper_ptr) {}
 

components/arc/test/fake_arc_bridge_host.h

@@ -75,6 +75,8 @@ class FakeArcBridgeHost : public mojom::ArcBridgeHost {
       mojom::VoiceInteractionFrameworkInstancePtr framework_ptr) override;
   void OnVolumeMounterInstanceReady(
       mojom::VolumeMounterInstancePtr volume_mounter_ptr) override;
+  void OnWakeLockInstanceReady(
+      mojom::WakeLockInstancePtr wakelock_ptr) override;
   void OnWallpaperInstanceReady(
       mojom::WallpaperInstancePtr wallpaper_ptr) override;
 

components/arc/test/fake_wake_lock_instance.cc

+// Copyright 2018 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include <utility>
+
+#include "components/arc/test/fake_wake_lock_instance.h"
+
+namespace arc {
+
+FakeWakeLockInstance::FakeWakeLockInstance() = default;
+
+FakeWakeLockInstance::~FakeWakeLockInstance() = default;
+
+void FakeWakeLockInstance::Init(mojom::WakeLockHostPtr host_ptr,
+                                InitCallback callback) {
+  host_ptr_ = std::move(host_ptr);
+  std::move(callback).Run();
+}
+
+}  // namespace arc

components/arc/test/fake_wake_lock_instance.h

+// Copyright 2018 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef COMPONENTS_ARC_TEST_FAKE_WAKE_LOCK_INSTANCE_H_
+#define COMPONENTS_ARC_TEST_FAKE_WAKE_LOCK_INSTANCE_H_
+
+#include "base/macros.h"
+#include "components/arc/common/wake_lock.mojom.h"
+
+namespace arc {
+
+class FakeWakeLockInstance : public mojom::WakeLockInstance {
+ public:
+  FakeWakeLockInstance();
+  ~FakeWakeLockInstance() override;
+
+  // mojom::WakeLockInstance overrides:
+  void Init(mojom::WakeLockHostPtr host_ptr, InitCallback callback) override;
+
+ private:
+  mojom::WakeLockHostPtr host_ptr_;
+
+  DISALLOW_COPY_AND_ASSIGN(FakeWakeLockInstance);
+};
+
+}  // namespace arc
+
+#endif  // COMPONENTS_ARC_TEST_FAKE_WAKE_LOCK_INSTANCE_H_

components/arc/wake_lock/DEPS

+include_rules = [
+  "+content/public/common/service_manager_connection.h",
+  "+services/device/public",
+  "+services/service_manager/public/cpp",
+]

components/arc/wake_lock/arc_wake_lock_bridge.cc

+// Copyright 2018 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include <memory>
+#include <utility>
+
+#include "base/logging.h"
+#include "base/memory/singleton.h"
+#include "chromeos/dbus/power_policy_controller.h"
+#include "components/arc/arc_bridge_service.h"
+#include "components/arc/arc_browser_context_keyed_service_factory_base.h"
+#include "components/arc/arc_service_manager.h"
+#include "components/arc/wake_lock/arc_wake_lock_bridge.h"
+#include "mojo/public/cpp/bindings/interface_request.h"
+#include "services/device/public/mojom/constants.mojom.h"
+#include "services/device/public/mojom/wake_lock_provider.mojom.h"
+#include "services/service_manager/public/cpp/connector.h"
+
+namespace arc {
+
+namespace {
+
+constexpr char kWakeLockReason[] = "ARC";
+
+// Singleton factory for ArcWakeLockBridge.
+class ArcWakeLockBridgeFactory
+    : public internal::ArcBrowserContextKeyedServiceFactoryBase<
+          ArcWakeLockBridge,
+          ArcWakeLockBridgeFactory> {
+ public:
+  // Factory name used by ArcBrowserContextKeyedServiceFactoryBase.
+  static constexpr const char* kName = "ArcWakeLockBridgeFactory";
+
+  static ArcWakeLockBridgeFactory* GetInstance() {
+    return base::Singleton<ArcWakeLockBridgeFactory>::get();
+  }
+
+ private:
+  friend base::DefaultSingletonTraits<ArcWakeLockBridgeFactory>;
+  ArcWakeLockBridgeFactory() = default;
+  ~ArcWakeLockBridgeFactory() override = default;
+};
+
+}  // namespace
+
+// WakeLockRequester requests a wake lock from the device service in response
+// to wake lock requests of a given type from Android. A count is kept of
+// outstanding Android requests so that only a single actual wake lock is used.
+class ArcWakeLockBridge::WakeLockRequester {
+ public:
+  WakeLockRequester(device::mojom::WakeLockType type,
+                    service_manager::Connector* connector)
+      : type_(type), connector_(connector) {}
+  ~WakeLockRequester() = default;
+
+  // Increments the number of outstanding requests from Android and requests a
+  // wake lock from the device service if this is the only request.
+  void AddRequest() {
+    DCHECK_GE(wake_lock_count_, 0);
+    wake_lock_count_++;
+    if (wake_lock_count_ > 1) {
+      DVLOG(1) << "Partial wake lock acquire. Count: " << wake_lock_count_;
+      return;
+    }
+
+    // Initialize |wake_lock_| if this is the first time we're using it.
+    DVLOG(1) << "Partial wake lock new acquire. Count: " << wake_lock_count_;
+    if (!wake_lock_) {
+      device::mojom::WakeLockProviderPtr provider;
+      connector_->BindInterface(device::mojom::kServiceName,
+                                mojo::MakeRequest(&provider));
+      provider->GetWakeLockWithoutContext(
+          type_, device::mojom::WakeLockReason::kOther, kWakeLockReason,
+          mojo::MakeRequest(&wake_lock_));
+    }
+
+    wake_lock_->RequestWakeLock();
+  }
+
+  // Decrements the number of outstanding Android requests. Cancels the device
+  // service wake lock when the request count hits zero.
+  void RemoveRequest() {
+    DCHECK_GE(wake_lock_count_, 0);
+    if (wake_lock_count_ == 0) {
+      LOG(WARNING) << "Release without acquire. Count: " << wake_lock_count_;
+      return;
+    }
+
+    wake_lock_count_--;
+    if (wake_lock_count_ >= 1) {
+      DVLOG(1) << "Partial wake release. Count: " << wake_lock_count_;
+      return;
+    }
+
+    DCHECK(wake_lock_);
+    DVLOG(1) << "Partial wake force release. Count: " << wake_lock_count_;
+    wake_lock_->CancelWakeLock();
+  }
+
+  // Runs the message loop until replies have been received for all pending
+  // requests on |wake_lock_|.
+  void FlushForTesting() {
+    if (wake_lock_)
+      wake_lock_.FlushForTesting();
+  }
+
+ private:
+  // Type of wake lock to request.
+  device::mojom::WakeLockType type_;
+
+  // Used to get services. Not owned.
+  service_manager::Connector* const connector_ = nullptr;
+
+  // Number of outstanding Android requests.
+  int64_t wake_lock_count_ = 0;
+
+  // Lazily initialized in response to first request.
+  device::mojom::WakeLockPtr wake_lock_;
+
+  DISALLOW_COPY_AND_ASSIGN(WakeLockRequester);
+};
+
+// static
+BrowserContextKeyedServiceFactory* ArcWakeLockBridge::GetFactory() {
+  return ArcWakeLockBridgeFactory::GetInstance();
+}
+
+// static
+ArcWakeLockBridge* ArcWakeLockBridge::GetForBrowserContext(
+    content::BrowserContext* context) {
+  return ArcWakeLockBridgeFactory::GetForBrowserContext(context);
+}
+
+// static
+ArcWakeLockBridge* ArcWakeLockBridge::GetForBrowserContextForTesting(
+    content::BrowserContext* context) {
+  return ArcWakeLockBridgeFactory::GetForBrowserContextForTesting(context);
+}
+
+ArcWakeLockBridge::ArcWakeLockBridge(content::BrowserContext* context,
+                                     ArcBridgeService* bridge_service)
+    : arc_bridge_service_(bridge_service),
+      binding_(this),
+      weak_ptr_factory_(this) {
+  arc_bridge_service_->wake_lock()->SetHost(this);
+  arc_bridge_service_->wake_lock()->AddObserver(this);
+}
+
+ArcWakeLockBridge::~ArcWakeLockBridge() {
+  arc_bridge_service_->wake_lock()->RemoveObserver(this);
+  arc_bridge_service_->wake_lock()->SetHost(nullptr);
+}
+
+void ArcWakeLockBridge::OnConnectionClosed() {
+  DVLOG(1) << "OnConnectionClosed";
+  wake_lock_requesters_.clear();
+}
+
+void ArcWakeLockBridge::FlushWakeLocksForTesting() {
+  for (const auto& it : wake_lock_requesters_)
+    it.second->FlushForTesting();
+}
+
+void ArcWakeLockBridge::AcquirePartialWakeLock(
+    AcquirePartialWakeLockCallback callback) {
+  GetWakeLockRequester(device::mojom::WakeLockType::kPreventAppSuspension)
+      ->AddRequest();
+  std::move(callback).Run(true);
+}
+
+void ArcWakeLockBridge::ReleasePartialWakeLock(
+    ReleasePartialWakeLockCallback callback) {
+  GetWakeLockRequester(device::mojom::WakeLockType::kPreventAppSuspension)
+      ->RemoveRequest();
+  std::move(callback).Run(true);
+}
+
+ArcWakeLockBridge::WakeLockRequester* ArcWakeLockBridge::GetWakeLockRequester(
+    device::mojom::WakeLockType type) {
+  auto it = wake_lock_requesters_.find(type);
+  if (it != wake_lock_requesters_.end())
+    return it->second.get();
+
+  service_manager::Connector* connector =
+      connector_for_test_
+          ? connector_for_test_
+          : content::ServiceManagerConnection::GetForProcess()->GetConnector();
+  DCHECK(connector);
+
+  it = wake_lock_requesters_
+           .emplace(type, std::make_unique<WakeLockRequester>(type, connector))
+           .first;
+  return it->second.get();
+}
+
+}  // namespace arc

components/arc/wake_lock/arc_wake_lock_bridge.h

+// Copyright 2018 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef COMPONENTS_ARC_WAKE_LOCK_ARC_WAKE_LOCK_BRIDGE_H_
+#define COMPONENTS_ARC_WAKE_LOCK_ARC_WAKE_LOCK_BRIDGE_H_
+
+#include <map>
+#include <memory>
+
+#include "base/macros.h"
+#include "base/memory/weak_ptr.h"
+#include "components/arc/common/wake_lock.mojom.h"
+#include "components/arc/connection_observer.h"
+#include "components/keyed_service/core/keyed_service.h"
+#include "content/public/common/service_manager_connection.h"
+#include "mojo/public/cpp/bindings/binding.h"
+#include "services/device/public/mojom/wake_lock.mojom.h"
+
+class BrowserContextKeyedServiceFactory;
+
+namespace content {
+class BrowserContext;
+}  // namespace content
+
+namespace arc {
+
+class ArcBridgeService;
+
+// Sets wake up timers / alarms based on calls from the instance.
+class ArcWakeLockBridge : public KeyedService,
+                          public ConnectionObserver<mojom::WakeLockInstance>,
+                          public mojom::WakeLockHost {
+ public:
+  // Returns the factory instance for this class.
+  static BrowserContextKeyedServiceFactory* GetFactory();
+
+  // Returns singleton instance for the given BrowserContext,
+  // or nullptr if the browser |context| is not allowed to use ARC.
+  static ArcWakeLockBridge* GetForBrowserContext(
+      content::BrowserContext* context);
+
+  static ArcWakeLockBridge* GetForBrowserContextForTesting(
+      content::BrowserContext* context);
+
+  ArcWakeLockBridge(content::BrowserContext* context,
+                    ArcBridgeService* bridge_service);
+  ~ArcWakeLockBridge() override;
+
+  void set_connector_for_testing(service_manager::Connector* connector) {
+    connector_for_test_ = connector;
+  }
+
+  // ConnectionObserver<mojom::WakeLockInstance>::Observer overrides.
+  void OnConnectionClosed() override;
+
+  // Runs the message loop until replies have been received for all pending
+  // device service requests in |wake_lock_requesters_|.
+  void FlushWakeLocksForTesting();
+
+  // mojom::WakeLockHost overrides.
+  void AcquirePartialWakeLock(AcquirePartialWakeLockCallback callback) override;
+  void ReleasePartialWakeLock(ReleasePartialWakeLockCallback callback) override;
+
+ private:
+  class WakeLockRequester;
+
+  // Returns the WakeLockRequester for |type|, creating one if needed.
+  WakeLockRequester* GetWakeLockRequester(device::mojom::WakeLockType type);
+
+  ArcBridgeService* const arc_bridge_service_;  // Owned by ArcServiceManager.
+
+  // If non-null, used instead of the process-wide connector to fetch services.
+  service_manager::Connector* connector_for_test_ = nullptr;
+
+  // Used to track Android wake lock requests and acquire and release device
+  // service wake locks as needed.
+  std::map<device::mojom::WakeLockType, std::unique_ptr<WakeLockRequester>>
+      wake_lock_requesters_;
+
+  mojo::Binding<mojom::WakeLockHost> binding_;
+
+  base::WeakPtrFactory<ArcWakeLockBridge> weak_ptr_factory_;
+
+  DISALLOW_COPY_AND_ASSIGN(ArcWakeLockBridge);
+};
+
+}  // namespace arc
+
+#endif  // COMPONENTS_ARC_WAKE_LOCK_ARC_WAKE_LOCK_BRIDGE_H_

components/arc/wake_lock/arc_wake_lock_bridge_unittest.cc

+// Copyright 2018 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "components/arc/wake_lock/arc_wake_lock_bridge.h"
+
+#include <memory>
+#include <utility>
+
+#include "base/run_loop.h"
+#include "base/test/scoped_task_environment.h"
+#include "components/arc/arc_bridge_service.h"
+#include "components/arc/common/power.mojom.h"
+#include "components/arc/test/connection_holder_util.h"
+#include "components/arc/test/fake_power_instance.h"
+#include "components/arc/test/fake_wake_lock_instance.h"
+#include "services/device/public/cpp/test/test_wake_lock_provider.h"
+#include "services/service_manager/public/cpp/test/test_connector_factory.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace arc {
+
+using device::mojom::WakeLockType;
+
+class ArcWakeLockBridgeTest : public testing::Test {
+ public:
+  ArcWakeLockBridgeTest() {
+    auto wake_lock_provider_ptr =
+        std::make_unique<device::TestWakeLockProvider>();
+    wake_lock_provider_ = wake_lock_provider_ptr.get();
+
+    connector_factory_ =
+        service_manager::TestConnectorFactory::CreateForUniqueService(
+            std::move(wake_lock_provider_ptr));
+    connector_ = connector_factory_->CreateConnector();
+
+    bridge_service_ = std::make_unique<ArcBridgeService>();
+    wake_lock_bridge_ =
+        std::make_unique<ArcWakeLockBridge>(nullptr, bridge_service_.get());
+    wake_lock_bridge_->set_connector_for_testing(connector_.get());
+    CreateWakeLockInstance();
+  }
+
+  ~ArcWakeLockBridgeTest() override { DestroyWakeLockInstance(); }
+
+ protected:
+  // Returns true iff there is no failure acquiring a system wake lock.
+  bool AcquirePartialWakeLock() {
+    base::RunLoop loop;
+    bool result = false;
+    wake_lock_bridge_->AcquirePartialWakeLock(base::BindOnce(
+        [](bool* result_out, bool result) { *result_out = result; }, &result));
+    loop.RunUntilIdle();
+    wake_lock_bridge_->FlushWakeLocksForTesting();
+    return result;
+  }
+
+  // Returns true iff there is no failure releasing a system wake lock.
+  bool ReleasePartialWakeLock() {
+    base::RunLoop loop;
+    bool result = false;
+    wake_lock_bridge_->ReleasePartialWakeLock(base::BindOnce(
+        [](bool* result_out, bool result) { *result_out = result; }, &result));
+    loop.RunUntilIdle();
+    wake_lock_bridge_->FlushWakeLocksForTesting();
+    return result;
+  }
+
+  // Creates a FakeWakeLockInstance for |bridge_service_|. This results in
+  // ArcWakeLockBridge::OnInstanceReady() being called.
+  void CreateWakeLockInstance() {
+    instance_ = std::make_unique<FakeWakeLockInstance>();
+    bridge_service_->wake_lock()->SetInstance(instance_.get());
+    WaitForInstanceReady(bridge_service_->wake_lock());
+  }
+
+  // Destroys the FakeWakeLockInstance. This results in
+  // ArcWakeLockBridge::OnInstanceClosed() being called.
+  void DestroyWakeLockInstance() {
+    if (!instance_)
+      return;
+    bridge_service_->wake_lock()->CloseInstance(instance_.get());
+    instance_.reset();
+  }
+
+  device::TestWakeLockProvider* GetWakeLockProvider() {
+    return wake_lock_provider_;
+  }
+
+ private:
+  base::test::ScopedTaskEnvironment scoped_task_environment_;
+
+  std::unique_ptr<service_manager::TestConnectorFactory> connector_factory_;
+  std::unique_ptr<service_manager::Connector> connector_;
+  device::TestWakeLockProvider* wake_lock_provider_;
+
+  std::unique_ptr<ArcBridgeService> bridge_service_;
+  std::unique_ptr<FakeWakeLockInstance> instance_;
+  std::unique_ptr<ArcWakeLockBridge> wake_lock_bridge_;
+
+  DISALLOW_COPY_AND_ASSIGN(ArcWakeLockBridgeTest);
+};
+
+TEST_F(ArcWakeLockBridgeTest, AcquireAndReleaseSinglePartialWakeLock) {
+  EXPECT_TRUE(AcquirePartialWakeLock());
+  EXPECT_EQ(1, GetWakeLockProvider()->GetActiveWakeLocksOfType(
+                   WakeLockType::kPreventAppSuspension));
+
+  EXPECT_TRUE(ReleasePartialWakeLock());
+  EXPECT_EQ(0, GetWakeLockProvider()->GetActiveWakeLocksOfType(
+                   WakeLockType::kPreventAppSuspension));
+}
+
+TEST_F(ArcWakeLockBridgeTest, AcquireAndReleaseMultiplePartialWakeLocks) {
+  // Taking multiple wake locks should result in only one active wake lock.
+  EXPECT_TRUE(AcquirePartialWakeLock());
+  EXPECT_TRUE(AcquirePartialWakeLock());
+  EXPECT_TRUE(AcquirePartialWakeLock());
+  EXPECT_EQ(1, GetWakeLockProvider()->GetActiveWakeLocksOfType(
+                   WakeLockType::kPreventAppSuspension));
+
+  // Releasing two wake locks after acquiring three should not result in
+  // releasing a wake lock.
+  EXPECT_TRUE(ReleasePartialWakeLock());
+  EXPECT_TRUE(ReleasePartialWakeLock());
+  EXPECT_EQ(1, GetWakeLockProvider()->GetActiveWakeLocksOfType(
+                   WakeLockType::kPreventAppSuspension));
+
+  // Releasing the remaining wake lock should result in the release of the wake
+  // lock.
+  EXPECT_TRUE(ReleasePartialWakeLock());
+  EXPECT_EQ(0, GetWakeLockProvider()->GetActiveWakeLocksOfType(
+                   WakeLockType::kPreventAppSuspension));
+}
+
+TEST_F(ArcWakeLockBridgeTest, ReleaseWakeLockOnInstanceClosed) {
+  EXPECT_TRUE(AcquirePartialWakeLock());
+  ASSERT_EQ(1, GetWakeLockProvider()->GetActiveWakeLocksOfType(
+                   WakeLockType::kPreventAppSuspension));
+
+  // If the instance is closed, all wake locks should be released.
+  base::RunLoop run_loop;
+  GetWakeLockProvider()->set_wake_lock_canceled_callback(
+      run_loop.QuitClosure());
+  DestroyWakeLockInstance();
+  run_loop.Run();
+  EXPECT_EQ(0, GetWakeLockProvider()->GetActiveWakeLocksOfType(
+                   WakeLockType::kPreventDisplaySleep));
+
+  // Check that wake locks can be requested after the instance becomes ready
+  // again.
+  CreateWakeLockInstance();
+  EXPECT_TRUE(AcquirePartialWakeLock());
+  EXPECT_EQ(1, GetWakeLockProvider()->GetActiveWakeLocksOfType(
+                   WakeLockType::kPreventAppSuspension));
+}
+
+}  // namespace arc