1 files changed, 51 insertions, 34 deletions

diff --git a/device/hid/hid_connection_unittest.cc b/device/hid/hid_connection_unittest.cc
index fbacb52..f4bf5de 100644
--- a/device/hid/hid_connection_unittest.cc
+++ b/device/hid/hid_connection_unittest.cc
@@ -9,6 +9,7 @@
 #include "base/callback.h"
 #include "base/memory/scoped_ptr.h"
 #include "base/run_loop.h"
+#include "base/scoped_observer.h"
 #include "device/hid/hid_connection.h"
 #include "device/hid/hid_service.h"
 #include "device/test/usb_test_gadget.h"
@@ -21,6 +22,54 @@ namespace {
 
 using net::IOBufferWithSize;
 
+// Helper class that can be used to block until a HID device with a particular
+// serial number is available. Example usage:
+//
+//   DeviceCatcher device_catcher("ABC123");
+//   HidDeviceId device_id = device_catcher.WaitForDevice();
+//   /* Call HidService::Connect(device_id) to open the device. */
+//
+class DeviceCatcher : HidService::Observer {
+ public:
+  DeviceCatcher(const std::string& serial_number)
+      : serial_number_(serial_number), observer_(this) {
+    HidService* hid_service = HidService::GetInstance(
+        base::MessageLoop::current()->message_loop_proxy());
+    observer_.Add(hid_service);
+    hid_service->GetDevices(base::Bind(&DeviceCatcher::OnEnumerationComplete,
+                                       base::Unretained(this)));
+  }
+
+  const HidDeviceId& WaitForDevice() {
+    run_loop_.Run();
+    observer_.RemoveAll();
+    return device_id_;
+  }
+
+ private:
+  void OnEnumerationComplete(const std::vector<HidDeviceInfo>& devices) {
+    for (const HidDeviceInfo& device_info : devices) {
+      if (device_info.serial_number == serial_number_) {
+        device_id_ = device_info.device_id;
+        run_loop_.Quit();
+        break;
+      }
+    }
+  }
+
+  void OnDeviceAdded(const HidDeviceInfo& device_info) override {
+    if (device_info.serial_number == serial_number_) {
+      device_id_ = device_info.device_id;
+      run_loop_.Quit();
+    }
+  }
+
+  std::string serial_number_;
+  ScopedObserver<device::HidService, device::HidService::Observer> observer_;
+  base::RunLoop run_loop_;
+  HidDeviceId device_id_;
+};
+
 class TestConnectCallback {
  public:
   TestConnectCallback()
@@ -105,43 +154,11 @@ class HidConnectionTest : public testing::Test {
     ASSERT_TRUE(test_gadget_);
     ASSERT_TRUE(test_gadget_->SetType(UsbTestGadget::HID_ECHO));
 
-    device_id_ = kInvalidHidDeviceId;
-
-    base::RunLoop run_loop;
-    message_loop_->PostDelayedTask(
-        FROM_HERE,
-        base::Bind(&HidConnectionTest::FindDevice,
-                   base::Unretained(this), run_loop.QuitClosure(), 5),
-        base::TimeDelta::FromMilliseconds(250));
-    run_loop.Run();
-
+    DeviceCatcher device_catcher(test_gadget_->GetSerialNumber());
+    device_id_ = device_catcher.WaitForDevice();
     ASSERT_NE(device_id_, kInvalidHidDeviceId);
   }
 
-  void FindDevice(const base::Closure& done, int retries) {
-    std::vector<HidDeviceInfo> devices;
-    service_->GetDevices(&devices);
-
-    for (std::vector<HidDeviceInfo>::iterator it = devices.begin();
-         it != devices.end();
-         ++it) {
-      if (it->serial_number == test_gadget_->GetSerialNumber()) {
-        device_id_ = it->device_id;
-        break;
-      }
-    }
-
-    if (device_id_ == kInvalidHidDeviceId && --retries > 0) {
-      message_loop_->PostDelayedTask(
-          FROM_HERE,
-          base::Bind(&HidConnectionTest::FindDevice, base::Unretained(this),
-                     done, retries),
-          base::TimeDelta::FromMilliseconds(10));
-    } else {
-      message_loop_->PostTask(FROM_HERE, done);
-    }
-  }
-
   scoped_ptr<base::MessageLoopForIO> message_loop_;
   HidService* service_;
   scoped_ptr<UsbTestGadget> test_gadget_;