Fix MDNS issue that was delaying/suppressing device listings.

Post-fix, MDNS device listings are returned instantly.
* Changed logic to compute the incremental change in active
  listeners by comparing service type histograms, rather than
  comparing the set of service types. The previous method was unable
  to detect the case when additional event handlers were registered
  for the same service type, so these event handlers did not
  promptly receive discovery data.
* Add "Refresh" method to DnsSdRegistry. I plan on exposing an
  extension API method to manually invoke this in the future.
* Made some parameters const refs in DnsSdRegistry.
* Encapsulate GetEventListenersByName and extension registry
  interactions with mockable functions for testability.
* Add tests and testability hooks for new behavior.

The bug here manifests in event pages because
EventListenerMap::LoadFilteredLazyListeners() competes for
initial event handler registration.

R=wez@chromium.org,apacible@chromium.org
BUG=504477,470369

Review URL: https://codereview.chromium.org/1192873004

Cr-Commit-Position: refs/heads/master@{#336430}

6 files changed, 364 insertions, 190 deletions

diff --git a/chrome/browser/extensions/api/mdns/dns_sd_registry.cc b/chrome/browser/extensions/api/mdns/dns_sd_registry.cc
index b8d7e80..7182c74 100644
--- a/chrome/browser/extensions/api/mdns/dns_sd_registry.cc
+++ b/chrome/browser/extensions/api/mdns/dns_sd_registry.cc
@@ -127,7 +127,11 @@ DnsSdDeviceLister* DnsSdRegistry::CreateDnsSdDeviceLister(
   return new DnsSdDeviceLister(discovery_client, delegate, service_type);
 }
 
-void DnsSdRegistry::RegisterDnsSdListener(std::string service_type) {
+void DnsSdRegistry::Refresh(const std::string& service_type) {
+  DispatchApiEvent(service_type);
+}
+
+void DnsSdRegistry::RegisterDnsSdListener(const std::string& service_type) {
   VLOG(1) << "RegisterDnsSdListener: " << service_type
           << ", registered: " << IsRegistered(service_type);
   if (service_type.empty())
@@ -148,7 +152,7 @@ void DnsSdRegistry::RegisterDnsSdListener(std::string service_type) {
   DispatchApiEvent(service_type);
 }
 
-void DnsSdRegistry::UnregisterDnsSdListener(std::string service_type) {
+void DnsSdRegistry::UnregisterDnsSdListener(const std::string& service_type) {
   VLOG(1) << "UnregisterDnsSdListener: " << service_type;
   DnsSdRegistry::DnsSdServiceTypeDataMap::iterator it =
       service_data_map_.find(service_type);
diff --git a/chrome/browser/extensions/api/mdns/dns_sd_registry.h b/chrome/browser/extensions/api/mdns/dns_sd_registry.h
index 2a97194..4af5344 100644
--- a/chrome/browser/extensions/api/mdns/dns_sd_registry.h
+++ b/chrome/browser/extensions/api/mdns/dns_sd_registry.h
@@ -43,13 +43,16 @@ class DnsSdRegistry : public DnsSdDelegate {
   explicit DnsSdRegistry(local_discovery::ServiceDiscoverySharedClient* client);
   virtual ~DnsSdRegistry();
 
+  // Broadcasts the device list for |service_type| to event listeners.
+  virtual void Refresh(const std::string& service_type);
+
   // Observer registration for parties interested in discovery events.
   virtual void AddObserver(DnsSdObserver* observer);
   virtual void RemoveObserver(DnsSdObserver* observer);
 
   // DNS-SD-related discovery functionality.
-  virtual void RegisterDnsSdListener(std::string service_type);
-  virtual void UnregisterDnsSdListener(std::string service_type);
+  virtual void RegisterDnsSdListener(const std::string& service_type);
+  virtual void UnregisterDnsSdListener(const std::string& service_type);
 
  protected:
   // Data class for managing all the resources and information related to a
diff --git a/chrome/browser/extensions/api/mdns/mdns_api.cc b/chrome/browser/extensions/api/mdns/mdns_api.cc
index c80d3b2..c7467a6b 100644
--- a/chrome/browser/extensions/api/mdns/mdns_api.cc
+++ b/chrome/browser/extensions/api/mdns/mdns_api.cc
@@ -79,38 +79,58 @@ DnsSdRegistry* MDnsAPI::dns_sd_registry() {
 
 void MDnsAPI::OnListenerAdded(const EventListenerInfo& details) {
   DCHECK(thread_checker_.CalledOnValidThread());
-  UpdateMDnsListeners(details);
+  UpdateMDnsListeners();
 }
 
 void MDnsAPI::OnListenerRemoved(const EventListenerInfo& details) {
   DCHECK(thread_checker_.CalledOnValidThread());
-  UpdateMDnsListeners(details);
+  UpdateMDnsListeners();
 }
 
-void MDnsAPI::UpdateMDnsListeners(const EventListenerInfo& details) {
+void MDnsAPI::UpdateMDnsListeners() {
   std::set<std::string> new_service_types;
+  ServiceTypeCounts current_service_counts;
   GetValidOnServiceListListeners(
       "" /* service_type_filter - blank = all services */,
-      nullptr /* extension_ids */,
-      &new_service_types);
-
-  // Find all the added and removed service types since last update.
-  std::set<std::string> added_service_types =
-      base::STLSetDifference<std::set<std::string> >(
-          new_service_types, service_types_);
-  std::set<std::string> removed_service_types =
-      base::STLSetDifference<std::set<std::string> >(
-          service_types_, new_service_types);
-
-  // Update the registry.
+      nullptr /* extension_ids */, &current_service_counts);
+
   DnsSdRegistry* registry = dns_sd_registry();
-  for (const auto& srv : added_service_types) {
-    registry->RegisterDnsSdListener(srv);
-  }
-  for (const auto& srv : removed_service_types) {
-    registry->UnregisterDnsSdListener(srv);
+
+  // Check if the counts of per-service-type event handlers has changed since
+  // the previous invocation, and take appropriate action if a change was
+  // detected.
+  //
+  // mDNS registration is performed for difference(cur, previous).
+  // mDNS unregistration is performed for difference(previous, cur).
+  // The mDNS device list is refreshed if the listener count has grown for
+  // a service type in union(cur, previous).
+  ServiceTypeCounts::iterator i_cur = current_service_counts.begin();
+  ServiceTypeCounts::iterator i_prev = prev_service_counts_.begin();
+  while (i_cur != current_service_counts.end() ||
+         i_prev != prev_service_counts_.end()) {
+    if (i_prev == prev_service_counts_.end() ||
+        (i_cur != current_service_counts.end() &&
+         i_cur->first < i_prev->first)) {
+      DVLOG(2) << "Registering listener for mDNS service " << i_cur->first;
+      registry->RegisterDnsSdListener(i_cur->first);
+      i_cur++;
+    } else if (i_cur == current_service_counts.end() ||
+               (i_prev != prev_service_counts_.end() &&
+                i_prev->first < i_cur->first)) {
+      DVLOG(2) << "Unregistering listener for mDNS service " << i_prev->first;
+      registry->UnregisterDnsSdListener(i_prev->first);
+      i_prev++;
+    } else {
+      if (i_cur->second > i_prev->second) {
+        DVLOG(2) << "Additional listeners added for mDNS service "
+                 << i_cur->first;
+        registry->Refresh(i_cur->first);
+      }
+      ++i_cur;
+      ++i_prev;
+    }
   }
-  service_types_ = new_service_types;
+  prev_service_counts_.swap(current_service_counts);
 }
 
 void MDnsAPI::OnDnsSdEvent(const std::string& service_type,
@@ -120,7 +140,7 @@ void MDnsAPI::OnDnsSdEvent(const std::string& service_type,
   std::vector<linked_ptr<mdns::MDnsService> > args;
   for (DnsSdRegistry::DnsSdServiceList::const_iterator it = services.begin();
        it != services.end(); ++it) {
-    if (static_cast<long>(args.size()) ==
+    if (static_cast<int>(args.size()) ==
         api::mdns::MAX_SERVICE_INSTANCES_PER_EVENT) {
       // TODO(reddaly): This is not the most meaningful way of notifying the
       // application that something bad happened.  It will go to the user's
@@ -158,13 +178,27 @@ void MDnsAPI::OnDnsSdEvent(const std::string& service_type,
   extensions::EventRouter::Get(browser_context_)->BroadcastEvent(event.Pass());
 }
 
+const extensions::EventListenerMap::ListenerList& MDnsAPI::GetEventListeners() {
+  return extensions::EventRouter::Get(browser_context_)
+      ->listeners()
+      .GetEventListenersByName(mdns::OnServiceList::kEventName);
+}
+
+bool MDnsAPI::IsMDnsAllowed(const std::string& extension_id,
+                            const std::string& service_type) const {
+  const extensions::Extension* extension =
+      ExtensionRegistry::Get(browser_context_)
+          ->enabled_extensions()
+          .GetByID(extension_id);
+  return (extension && (extension->is_platform_app() ||
+                        IsServiceTypeWhitelisted(service_type)));
+}
+
 void MDnsAPI::GetValidOnServiceListListeners(
     const std::string& service_type_filter,
     std::set<std::string>* extension_ids,
-    std::set<std::string>* service_types) {
-  for (const auto& listener :
-         extensions::EventRouter::Get(browser_context_)->listeners()
-         .GetEventListenersByName(mdns::OnServiceList::kEventName)) {
+    ServiceTypeCounts* service_type_counts) {
+  for (const auto& listener : GetEventListeners()) {
     base::DictionaryValue* filter = listener->filter();
 
     std::string service_type;
@@ -176,22 +210,16 @@ void MDnsAPI::GetValidOnServiceListListeners(
     if (!service_type_filter.empty() && service_type_filter != service_type)
       continue;
 
-    const Extension* extension = ExtensionRegistry::Get(browser_context_)->
-        enabled_extensions().GetByID(listener->extension_id());
-    // Don't listen for services associated only with disabled extensions.
-    if (!extension)
-      continue;
-
-    // Platform apps may query for all services; other types of extensions are
-    // restricted to a whitelist.
-    if (!extension->is_platform_app() &&
-        !IsServiceTypeWhitelisted(service_type))
+    // Don't listen for services associated only with disabled extensions
+    // or non-whitelisted, non-platform-app extensions.
+    if (!IsMDnsAllowed(listener->extension_id(), service_type))
       continue;
 
     if (extension_ids)
       extension_ids->insert(listener->extension_id());
-    if (service_types)
-      service_types->insert(service_type);
+    if (service_type_counts) {
+      (*service_type_counts)[service_type]++;
+    }
   }
 }
 
@@ -201,9 +229,9 @@ void MDnsAPI::WriteToConsole(const std::string& service_type,
   // Get all the extensions with an onServiceList listener for a particular
   // service type.
   std::set<std::string> extension_ids;
-  GetValidOnServiceListListeners(service_type,
-                                 &extension_ids,
-                                 nullptr /* service_types */);
+  ServiceTypeCounts counts;
+  GetValidOnServiceListListeners(service_type, &extension_ids,
+                                 nullptr /* service_type_counts */);
 
   std::string logged_message(std::string("[chrome.mdns] ") + message);
 
diff --git a/chrome/browser/extensions/api/mdns/mdns_api.h b/chrome/browser/extensions/api/mdns/mdns_api.h
index 95e8428..c79bb1c 100644
--- a/chrome/browser/extensions/api/mdns/mdns_api.h
+++ b/chrome/browser/extensions/api/mdns/mdns_api.h
@@ -5,9 +5,11 @@
 #ifndef CHROME_BROWSER_EXTENSIONS_API_MDNS_MDNS_API_H_
 #define CHROME_BROWSER_EXTENSIONS_API_MDNS_MDNS_API_H_
 
+#include <map>
 #include <set>
 #include <string>
 
+#include "base/gtest_prod_util.h"
 #include "base/memory/scoped_ptr.h"
 #include "base/threading/thread_checker.h"
 #include "chrome/browser/extensions/api/mdns/dns_sd_registry.h"
@@ -46,7 +48,14 @@ class MDnsAPI : public BrowserContextKeyedAPI,
   // Retrieve an instance of the registry. Lazily created when needed.
   virtual DnsSdRegistry* dns_sd_registry();
 
+  // Gets the list of mDNS event listeners.
+  virtual const extensions::EventListenerMap::ListenerList& GetEventListeners();
+
  private:
+  FRIEND_TEST_ALL_PREFIXES(MDnsAPIDiscoveryTest,
+                           ServiceListenersAddedAndRemoved);
+
+  typedef std::map<std::string, int> ServiceTypeCounts;
   friend class BrowserContextKeyedAPIFactory<MDnsAPI>;
 
   // EventRouter::Observer:
@@ -66,7 +75,7 @@ class MDnsAPI : public BrowserContextKeyedAPI,
   static const bool kServiceIsNULLWhileTesting = true;
 
   // Update the current list of service types and update the registry.
-  void UpdateMDnsListeners(const EventListenerInfo& details);
+  void UpdateMDnsListeners();
 
   // Write a message to the consoles of extensions listening to a given service
   // type.
@@ -74,21 +83,29 @@ class MDnsAPI : public BrowserContextKeyedAPI,
                       content::ConsoleMessageLevel level,
                       const std::string& message);
 
+  // Returns true if an extension or platform app |extension_id| is allowed to
+  // listen to mDNS events for |service_type|.
+  virtual bool IsMDnsAllowed(const std::string& extension_id,
+                             const std::string& service_type) const;
+
   // Finds all all the valid listeners of the mdns.onServiceList event and
-  // filters them by service type if |service_type_filter| is non-empty.  The
-  // extension ids and matched service types are output to |extension_ids| and
-  // |service_types|, respectively, if the supplied pointers is non-null.
+  // filters them by service type if |service_type_filter| is non-empty.
+  // The list of extensions with active listeners is written to |extension_ids|,
+  // if non-null.
+  // The counts for each service type are output to |service_type_counts|, if
+  // non-null.
   void GetValidOnServiceListListeners(const std::string& service_type_filter,
                                       std::set<std::string>* extension_ids,
-                                      std::set<std::string>* service_types);
+                                      ServiceTypeCounts* service_type_counts);
 
   // Ensure methods are only called on UI thread.
   base::ThreadChecker thread_checker_;
   content::BrowserContext* const browser_context_;
   // Lazily created on first access and destroyed with this API class.
   scoped_ptr<DnsSdRegistry> dns_sd_registry_;
-  // Current set of service types registered with the registry.
-  std::set<std::string> service_types_;
+  // Count of active listeners per service type, saved from the previous
+  // invocation of UpdateMDnsListeners().
+  ServiceTypeCounts prev_service_counts_;
 
   DISALLOW_COPY_AND_ASSIGN(MDnsAPI);
 };
diff --git a/chrome/browser/extensions/api/mdns/mdns_api_unittest.cc b/chrome/browser/extensions/api/mdns/mdns_api_unittest.cc
index 7c8a190..31749c9 100644
--- a/chrome/browser/extensions/api/mdns/mdns_api_unittest.cc
+++ b/chrome/browser/extensions/api/mdns/mdns_api_unittest.cc
@@ -2,6 +2,9 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+#include <vector>
+
+#include "base/memory/linked_ptr.h"
 #include "base/memory/scoped_ptr.h"
 #include "chrome/browser/extensions/api/mdns/mdns_api.h"
 #include "chrome/browser/extensions/extension_service.h"
@@ -10,6 +13,7 @@
 #include "chrome/common/extensions/api/mdns.h"
 #include "content/public/browser/browser_context.h"
 #include "content/public/test/mock_render_process_host.h"
+#include "extensions/browser/event_listener_map.h"
 #include "extensions/browser/event_router_factory.h"
 #include "extensions/browser/extension_prefs.h"
 #include "extensions/browser/extension_registry.h"
@@ -17,10 +21,55 @@
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
-namespace extensions {
+using testing::_;
+using testing::Return;
+using testing::ReturnRef;
 
+namespace extensions {
 namespace {
 
+const char kExtId[] = "mbflcebpggnecokmikipoihdbecnjfoj";
+const char kService1[] = "service1";
+const char kService2[] = "service2";
+
+// Registers a new EventListener for |service_types| in |listener_list|.
+void AddEventListener(
+    const std::string& extension_id,
+    const std::string& service_type,
+    extensions::EventListenerMap::ListenerList* listener_list) {
+  scoped_ptr<base::DictionaryValue> filter(new base::DictionaryValue);
+  filter->SetString(kEventFilterServiceTypeKey, service_type);
+  listener_list->push_back(make_linked_ptr(
+      EventListener::ForExtension(kEventFilterServiceTypeKey, extension_id,
+                                  nullptr, filter.Pass()).release()));
+}
+
+class NullDelegate : public EventListenerMap::Delegate {
+ public:
+  void OnListenerAdded(const EventListener* listener) override {}
+  void OnListenerRemoved(const EventListener* listener) override {}
+};
+
+// Testing subclass of MDnsAPI which replaces calls to core browser components
+// that we don't want to have to instantiate or mock for these tests.
+class MockedMDnsAPI : public MDnsAPI {
+ public:
+  explicit MockedMDnsAPI(content::BrowserContext* context) : MDnsAPI(context) {}
+
+ public:
+  MOCK_CONST_METHOD2(IsMDnsAllowed,
+                     bool(const std::string& extension_id,
+                          const std::string& service_type));
+
+  MOCK_METHOD0(GetEventListeners,
+               const extensions::EventListenerMap::ListenerList&());
+};
+
+scoped_ptr<KeyedService> MockedMDnsAPITestingFactoryFunction(
+    content::BrowserContext* context) {
+  return make_scoped_ptr(new MockedMDnsAPI(context));
+}
+
 scoped_ptr<KeyedService> MDnsAPITestingFactoryFunction(
     content::BrowserContext* context) {
   return make_scoped_ptr(new MDnsAPI(context));
@@ -44,8 +93,9 @@ class MockDnsSdRegistry : public DnsSdRegistry {
 
   MOCK_METHOD1(AddObserver, void(DnsSdObserver* observer));
   MOCK_METHOD1(RemoveObserver, void(DnsSdObserver* observer));
-  MOCK_METHOD1(RegisterDnsSdListener, void(std::string service_type));
-  MOCK_METHOD1(UnregisterDnsSdListener, void(std::string service_type));
+  MOCK_METHOD1(RegisterDnsSdListener, void(const std::string& service_type));
+  MOCK_METHOD1(UnregisterDnsSdListener, void(const std::string& service_type));
+  MOCK_METHOD1(Refresh, void(const std::string& service_type));
 
   void DispatchMDnsEvent(const std::string& service_type,
                          const DnsSdServiceList& services) {
@@ -56,6 +106,78 @@ class MockDnsSdRegistry : public DnsSdRegistry {
   extensions::DnsSdRegistry::DnsSdObserver* api_;
 };
 
+class MockEventRouter : public EventRouter {
+ public:
+  explicit MockEventRouter(content::BrowserContext* browser_context,
+                           ExtensionPrefs* extension_prefs)
+      : EventRouter(browser_context, extension_prefs) {}
+  virtual ~MockEventRouter() {}
+
+  virtual void BroadcastEvent(scoped_ptr<Event> event) {
+    BroadcastEventPtr(event.get());
+  }
+  MOCK_METHOD1(BroadcastEventPtr, void(Event* event));
+};
+
+scoped_ptr<KeyedService> MockEventRouterFactoryFunction(
+    content::BrowserContext* context) {
+  return make_scoped_ptr(
+      new MockEventRouter(context, ExtensionPrefs::Get(context)));
+}
+
+class EventServiceListSizeMatcher
+    : public testing::MatcherInterface<const Event&> {
+ public:
+  explicit EventServiceListSizeMatcher(size_t expected_size)
+      : expected_size_(expected_size) {}
+
+  virtual bool MatchAndExplain(const Event& e,
+                               testing::MatchResultListener* listener) const {
+    if (!e.event_args) {
+      *listener << "event.event_arg is null when it shouldn't be";
+      return false;
+    }
+    if (e.event_args->GetSize() != 1) {
+      *listener << "event.event_arg.GetSize() should be 1 but is "
+                << e.event_args->GetSize();
+      return false;
+    }
+    const base::ListValue* services = nullptr;
+    {
+      const base::Value* out;
+      e.event_args->Get(0, &out);
+      services = static_cast<const base::ListValue*>(out);
+    }
+    if (!services) {
+      *listener << "event's service list argument is not a ListValue";
+      return false;
+    }
+    *listener << "number of services is " << services->GetSize();
+    return static_cast<testing::Matcher<size_t>>(testing::Eq(expected_size_))
+        .MatchAndExplain(services->GetSize(), listener);
+  }
+
+  virtual void DescribeTo(::std::ostream* os) const {
+    *os << "is an onServiceList event where the number of services is "
+        << expected_size_;
+  }
+
+  virtual void DescribeNegationTo(::std::ostream* os) const {
+    *os << "isn't an onServiceList event where the number of services is "
+        << expected_size_;
+  }
+
+ private:
+  size_t expected_size_;
+};
+
+inline testing::Matcher<const Event&> EventServiceListSize(
+    size_t expected_size) {
+  return testing::MakeMatcher(new EventServiceListSizeMatcher(expected_size));
+}
+
+}  // namespace
+
 class MDnsAPITest : public extensions::ExtensionServiceTestBase {
  public:
   void SetUp() override {
@@ -66,15 +188,14 @@ class MDnsAPITest : public extensions::ExtensionServiceTestBase {
 
     // A custom TestingFactoryFunction is required for an MDnsAPI to actually be
     // constructed.
-    MDnsAPI::GetFactoryInstance()->SetTestingFactory(
-        browser_context(),
-        MDnsAPITestingFactoryFunction);
+    MDnsAPI::GetFactoryInstance()->SetTestingFactory(browser_context(),
+                                                     GetMDnsFactory());
 
     EventRouterFactory::GetInstance()->SetTestingFactory(browser_context(),
                                                          &BuildEventRouter);
 
     // Do some sanity checking
-    ASSERT_TRUE(MDnsAPI::Get(browser_context())); // constructs MDnsAPI
+    ASSERT_TRUE(MDnsAPI::Get(browser_context()));      // constructs MDnsAPI
     ASSERT_TRUE(EventRouter::Get(browser_context()));  // constructs EventRouter
 
     registry_ = new MockDnsSdRegistry(MDnsAPI::Get(browser_context()));
@@ -88,6 +209,12 @@ class MDnsAPITest : public extensions::ExtensionServiceTestBase {
         new content::MockRenderProcessHost(browser_context()));
   }
 
+  // Returns the mDNS API factory function (mock vs. real) to use for the test.
+  virtual BrowserContextKeyedServiceFactory::TestingFactoryFunction
+  GetMDnsFactory() {
+    return MDnsAPITestingFactoryFunction;
+  }
+
   void TearDown() override {
     EXPECT_CALL(*dns_sd_registry(),
                 RemoveObserver(MDnsAPI::Get(browser_context())))
@@ -136,13 +263,107 @@ class MDnsAPITest : public extensions::ExtensionServiceTestBase {
   MockDnsSdRegistry* registry_;
 
   scoped_ptr<content::RenderProcessHost> render_process_host_;
+};
 
+class MDnsAPIMaxServicesTest : public MDnsAPITest {
+ public:
+  MockEventRouter* event_router() {
+    return static_cast<MockEventRouter*>(EventRouter::Get(browser_context()));
+  }
 };
 
+class MDnsAPIDiscoveryTest : public MDnsAPITest {
+ public:
+  BrowserContextKeyedServiceFactory::TestingFactoryFunction GetMDnsFactory()
+      override {
+    return MockedMDnsAPITestingFactoryFunction;
+  }
+
+  void SetUp() override {
+    MDnsAPITest::SetUp();
+    mdns_api_ = static_cast<MockedMDnsAPI*>(MDnsAPI::Get(browser_context()));
+    EXPECT_CALL(*mdns_api_, IsMDnsAllowed(_, _)).WillRepeatedly(Return(true));
+  }
+
+ protected:
+  MockedMDnsAPI* mdns_api_;
+};
+
+TEST_F(MDnsAPIDiscoveryTest, ServiceListenersAddedAndRemoved) {
+  EventRouterFactory::GetInstance()->SetTestingFactory(
+      browser_context(), &MockEventRouterFactoryFunction);
+  extensions::EventListenerMap::ListenerList listeners;
+
+  extensions::EventListenerInfo listener_info(
+      kEventFilterServiceTypeKey, kExtId, GURL(), browser_context());
+
+  EXPECT_CALL(*mdns_api_, GetEventListeners())
+      .WillRepeatedly(ReturnRef(listeners));
+
+  // Listener #1 added with kService1.
+  AddEventListener(kExtId, kService1, &listeners);
+  EXPECT_CALL(*dns_sd_registry(), RegisterDnsSdListener(kService1));
+  mdns_api_->OnListenerAdded(listener_info);
+
+  // Listener #2 added with kService2.
+  AddEventListener(kExtId, kService2, &listeners);
+  EXPECT_CALL(*dns_sd_registry(), RegisterDnsSdListener(kService2));
+  mdns_api_->OnListenerAdded(listener_info);
+
+  // No-op.
+  mdns_api_->OnListenerAdded(listener_info);
+
+  // Listener #3 added with kService2. Should trigger a refresh of kService2.
+  AddEventListener(kExtId, kService2, &listeners);
+  EXPECT_CALL(*dns_sd_registry(), Refresh(kService2));
+  mdns_api_->OnListenerAdded(listener_info);
+
+  // Listener #3 removed, should be a no-op since there is still a live
+  // listener for kService2.
+  listeners.pop_back();
+  mdns_api_->OnListenerRemoved(listener_info);
+
+  // Listener #2 removed, should unregister kService2.
+  listeners.pop_back();
+  EXPECT_CALL(*dns_sd_registry(), UnregisterDnsSdListener(kService2));
+  mdns_api_->OnListenerRemoved(listener_info);
+
+  // Listener #1 removed, should unregister kService1.
+  listeners.pop_back();
+  EXPECT_CALL(*dns_sd_registry(), UnregisterDnsSdListener(kService1));
+  mdns_api_->OnListenerRemoved(listener_info);
+
+  // No-op.
+  mdns_api_->OnListenerAdded(listener_info);
+
+  // Listener #4 added with kService1.
+  AddEventListener(kExtId, kService1, &listeners);
+  EXPECT_CALL(*dns_sd_registry(), RegisterDnsSdListener(kService1));
+  mdns_api_->OnListenerAdded(listener_info);
+}
+
+TEST_F(MDnsAPIMaxServicesTest, OnServiceListDoesNotExceedLimit) {
+  EventRouterFactory::GetInstance()->SetTestingFactory(
+      browser_context(), &MockEventRouterFactoryFunction);
+
+  // This check should change when the [value=64] changes in the IDL file.
+  EXPECT_EQ(64, api::mdns::MAX_SERVICE_INSTANCES_PER_EVENT);
+
+  // Dispatch an mDNS event with more service instances than the max, and ensure
+  // that the list is truncated by inspecting the argument to MockEventRouter's
+  // BroadcastEvent method.
+  DnsSdRegistry::DnsSdServiceList services;
+  for (int i = 0; i < api::mdns::MAX_SERVICE_INSTANCES_PER_EVENT + 10; ++i) {
+    services.push_back(DnsSdService());
+  }
+  EXPECT_CALL(*event_router(), BroadcastEventPtr(testing::Pointee(
+                                   EventServiceListSize(size_t(64))))).Times(1);
+  dns_sd_registry()->DispatchMDnsEvent("_testing._tcp.local", services);
+}
+
 TEST_F(MDnsAPITest, ExtensionRespectsWhitelist) {
-  const std::string ext_id("mbflcebpggnecokmikipoihdbecnjfoj");
   scoped_refptr<extensions::Extension> extension =
-      CreateExtension("Dinosaur networker", false, ext_id);
+      CreateExtension("Dinosaur networker", false, kExtId);
   ExtensionRegistry::Get(browser_context())->AddEnabled(extension);
   ASSERT_EQ(Manifest::TYPE_EXTENSION, extension.get()->GetType());
 
@@ -156,15 +377,17 @@ TEST_F(MDnsAPITest, ExtensionRespectsWhitelist) {
     // Test that the extension is able to listen to a non-whitelisted service
     EXPECT_CALL(*dns_sd_registry(), RegisterDnsSdListener("_trex._tcp.local"))
         .Times(0);
-    EventRouter::Get(browser_context())->AddFilteredEventListener(
-        api::mdns::OnServiceList::kEventName, render_process_host(), ext_id,
-        filter, false);
+    EventRouter::Get(browser_context())
+        ->AddFilteredEventListener(api::mdns::OnServiceList::kEventName,
+                                   render_process_host(), kExtId, filter,
+                                   false);
 
     EXPECT_CALL(*dns_sd_registry(), UnregisterDnsSdListener("_trex._tcp.local"))
         .Times(0);
-    EventRouter::Get(browser_context())->RemoveFilteredEventListener(
-        api::mdns::OnServiceList::kEventName, render_process_host(), ext_id,
-        filter, false);
+    EventRouter::Get(browser_context())
+        ->RemoveFilteredEventListener(api::mdns::OnServiceList::kEventName,
+                                      render_process_host(), kExtId, filter,
+                                      false);
   }
   {
     base::DictionaryValue filter;
@@ -174,22 +397,23 @@ TEST_F(MDnsAPITest, ExtensionRespectsWhitelist) {
     // Test that the extension is able to listen to a whitelisted service
     EXPECT_CALL(*dns_sd_registry(),
                 RegisterDnsSdListener("_testing._tcp.local"));
-    EventRouter::Get(browser_context())->AddFilteredEventListener(
-        api::mdns::OnServiceList::kEventName, render_process_host(), ext_id,
-        filter, false);
+    EventRouter::Get(browser_context())
+        ->AddFilteredEventListener(api::mdns::OnServiceList::kEventName,
+                                   render_process_host(), kExtId, filter,
+                                   false);
 
     EXPECT_CALL(*dns_sd_registry(),
                 UnregisterDnsSdListener("_testing._tcp.local"));
-    EventRouter::Get(browser_context())->RemoveFilteredEventListener(
-        api::mdns::OnServiceList::kEventName, render_process_host(),
-        ext_id, filter, false);
+    EventRouter::Get(browser_context())
+        ->RemoveFilteredEventListener(api::mdns::OnServiceList::kEventName,
+                                      render_process_host(), kExtId, filter,
+                                      false);
   }
 }
 
 TEST_F(MDnsAPITest, PlatformAppsNotSubjectToWhitelist) {
-  const std::string ext_id("mbflcebpggnecokmikipoihdbecnjfoj");
   scoped_refptr<extensions::Extension> extension =
-      CreateExtension("Dinosaur networker", true, ext_id);
+      CreateExtension("Dinosaur networker", true, kExtId);
   ExtensionRegistry::Get(browser_context())->AddEnabled(extension);
   ASSERT_TRUE(extension.get()->is_platform_app());
 
@@ -199,115 +423,15 @@ TEST_F(MDnsAPITest, PlatformAppsNotSubjectToWhitelist) {
   ASSERT_TRUE(dns_sd_registry());
   // Test that the extension is able to listen to a non-whitelisted service
   EXPECT_CALL(*dns_sd_registry(), RegisterDnsSdListener("_trex._tcp.local"));
-  EventRouter::Get(browser_context())->AddFilteredEventListener(
-      api::mdns::OnServiceList::kEventName, render_process_host(), ext_id,
-      filter, false);
+  EventRouter::Get(browser_context())
+      ->AddFilteredEventListener(api::mdns::OnServiceList::kEventName,
+                                 render_process_host(), kExtId, filter, false);
 
   EXPECT_CALL(*dns_sd_registry(), UnregisterDnsSdListener("_trex._tcp.local"));
-  EventRouter::Get(browser_context())->RemoveFilteredEventListener(
-      api::mdns::OnServiceList::kEventName, render_process_host(), ext_id,
-      filter, false);
-}
-
-class MockEventRouter : public EventRouter {
- public:
-  explicit MockEventRouter(content::BrowserContext* browser_context,
-                           ExtensionPrefs* extension_prefs) :
-      EventRouter(browser_context, extension_prefs) {}
-  virtual ~MockEventRouter() {}
-
-  virtual void BroadcastEvent(scoped_ptr<Event> event) {
-    BroadcastEventPtr(event.get());
-  }
-  MOCK_METHOD1(BroadcastEventPtr, void(Event* event));
-};
-
-scoped_ptr<KeyedService> MockEventRouterFactoryFunction(
-    content::BrowserContext* context) {
-  return make_scoped_ptr(
-      new MockEventRouter(context, ExtensionPrefs::Get(context)));
-}
-
-class MDnsAPIMaxServicesTest : public MDnsAPITest {
- public:
-  MockEventRouter* event_router() {
-    return static_cast<MockEventRouter*>(EventRouter::Get(browser_context()));
-  }
-};
-
-class EventServiceListSizeMatcher :
-      public testing::MatcherInterface<const Event&> {
- public:
-  explicit EventServiceListSizeMatcher(size_t expected_size)
-      : expected_size_(expected_size) {}
-
-  virtual bool MatchAndExplain(const Event& e,
-                               testing::MatchResultListener* listener) const {
-    if (e.event_args.get() == nullptr) {
-      *listener << "event.event_arg is null when it shouldn't be";
-      return false;
-    }
-    if (e.event_args->GetSize() != 1) {
-      *listener << "event.event_arg.GetSize() should be 1 but is "
-                << e.event_args->GetSize();
-      return false;
-    }
-    const base::ListValue* services = nullptr;
-    {
-      const base::Value* out;
-      e.event_args->Get(0, &out);
-      services = static_cast<const base::ListValue*>(out);
-    }
-    if (services == nullptr) {
-      *listener << "event's service list argument is not a ListValue";
-      return false;
-    }
-    *listener << "number of services is "
-              << services->GetSize();
-    return static_cast<testing::Matcher<size_t>>(testing::Eq(expected_size_))
-        .MatchAndExplain(services->GetSize(), listener);
-  }
-
-  virtual void DescribeTo(::std::ostream* os) const {
-    *os << "is an onServiceList event where the number of services is "
-        << expected_size_;
-  }
-
-  virtual void DescribeNegationTo(::std::ostream* os) const {
-    *os << "isn't an onServiceList event where the number of services is "
-        << expected_size_;
-  }
- private:
-  size_t expected_size_;
-};
-
-inline testing::Matcher<const Event&> EventServiceListSize(
-    size_t expected_size) {
-  return testing::MakeMatcher(new EventServiceListSizeMatcher(expected_size));
+  EventRouter::Get(browser_context())
+      ->RemoveFilteredEventListener(api::mdns::OnServiceList::kEventName,
+                                    render_process_host(), kExtId, filter,
+                                    false);
 }
 
-TEST_F(MDnsAPIMaxServicesTest, OnServiceListDoesNotExceedLimit) {
-  EventRouterFactory::GetInstance()->SetTestingFactory(
-      browser_context(), &MockEventRouterFactoryFunction);
-
-   // This check should change when the [value=64] changes in the IDL file.
-  EXPECT_EQ(64, api::mdns::MAX_SERVICE_INSTANCES_PER_EVENT);
-
-  // Dispatch an mDNS event with more service instances than the max, and ensure
-  // that the list is truncated by inspecting the argument to MockEventRouter's
-  // BroadcastEvent method.
-  DnsSdRegistry::DnsSdServiceList services;
-  for (int i=0; i < api::mdns::MAX_SERVICE_INSTANCES_PER_EVENT + 10; ++i) {
-    services.push_back(DnsSdService());
-  }
-  EXPECT_CALL(
-      *event_router(),
-      BroadcastEventPtr(
-          testing::Pointee(EventServiceListSize(size_t(64)))))
-      .Times(1);
-  dns_sd_registry()->DispatchMDnsEvent("_testing._tcp.local", services);
-}
-
-} // empty namespace
-
-} // namespace extensions
+}  // namespace extensions
diff --git a/chrome/browser/extensions/api/mdns/mdns_apitest.cc b/chrome/browser/extensions/api/mdns/mdns_apitest.cc
index f7b87be..b084a2a 100644
--- a/chrome/browser/extensions/api/mdns/mdns_apitest.cc
+++ b/chrome/browser/extensions/api/mdns/mdns_apitest.cc
@@ -12,8 +12,8 @@
 #include "testing/gmock/include/gmock/gmock.h"
 
 using extensions::DnsSdRegistry;
-using ::testing::A;
 using ::testing::_;
+using ::testing::A;
 
 namespace api = extensions::api;
 
@@ -26,8 +26,8 @@ class MockDnsSdRegistry : public DnsSdRegistry {
 
   MOCK_METHOD1(AddObserver, void(DnsSdObserver* observer));
   MOCK_METHOD1(RemoveObserver, void(DnsSdObserver* observer));
-  MOCK_METHOD1(RegisterDnsSdListener, void(std::string service_type));
-  MOCK_METHOD1(UnregisterDnsSdListener, void(std::string service_type));
+  MOCK_METHOD1(RegisterDnsSdListener, void(const std::string& service_type));
+  MOCK_METHOD1(UnregisterDnsSdListener, void(const std::string& service_type));
 
   void DispatchMDnsEvent(const std::string& service_type,
                          const DnsSdServiceList& services) {
@@ -154,10 +154,8 @@ IN_PROC_BROWSER_TEST_F(MDnsAPITest, MAYBE_RegisterMultipleListeners) {
 IN_PROC_BROWSER_TEST_F(MDnsAPITest, MAYBE_RegisterTooManyListeners) {
   SetUpTestDnsSdRegistry();
 
-  EXPECT_CALL(*dns_sd_registry_, RegisterDnsSdListener(A<std::string>()))
-      .Times(10);
-  EXPECT_CALL(*dns_sd_registry_, UnregisterDnsSdListener(A<std::string>()))
-      .Times(10);
+  EXPECT_CALL(*dns_sd_registry_, RegisterDnsSdListener(_)).Times(10);
+  EXPECT_CALL(*dns_sd_registry_, UnregisterDnsSdListener(_)).Times(10);
   EXPECT_CALL(*dns_sd_registry_,
               RemoveObserver(A<extensions::DnsSdRegistry::DnsSdObserver*>()))
       .Times(1);