5 files changed, 99 insertions, 3 deletions
diff --git a/mojo/mojo.gyp b/mojo/mojo.gyp
index 6e3acbf..53e7324 100644
--- a/mojo/mojo.gyp
+++ b/mojo/mojo.gyp
@@ -45,6 +45,7 @@
         'mojo_network_service',
         'mojo_pepper_container_app',
         'mojo_profile_service',
+        'mojo_public_application_unittests',
         'mojo_public_test_utils',
         'mojo_public_bindings_unittests',
         'mojo_public_environment_unittests',
diff --git a/mojo/mojo_public.gypi b/mojo/mojo_public.gypi
index 2690c7c..77b755d 100644
--- a/mojo/mojo_public.gypi
+++ b/mojo/mojo_public.gypi
@@ -166,6 +166,21 @@
       ],
     },
     {
+      'target_name': 'mojo_public_application_unittests',
+      'type': 'executable',
+      'dependencies': [
+        '../base/base.gyp:base',
+        '../testing/gtest.gyp:gtest',
+        'mojo_application',
+        'mojo_run_all_unittests',
+        'mojo_environment_standalone',
+        'mojo_utility',
+      ],
+      'sources': [
+        'public/cpp/application/tests/service_registry_unittest.cc',
+      ],
+    },
+    {
       'target_name': 'mojo_public_system_unittests',
       'type': 'executable',
       'dependencies': [
diff --git a/mojo/public/cpp/application/lib/service_registry.cc b/mojo/public/cpp/application/lib/service_registry.cc
index bc901df..44b51ef 100644
--- a/mojo/public/cpp/application/lib/service_registry.cc
+++ b/mojo/public/cpp/application/lib/service_registry.cc
@@ -33,19 +33,27 @@ ServiceRegistry::~ServiceRegistry() {
 
 void ServiceRegistry::AddServiceConnector(
     ServiceConnectorBase* service_connector) {
+  RemoveServiceConnectorInternal(service_connector);
   name_to_service_connector_[service_connector->name()] = service_connector;
   service_connector->set_registry(this);
 }
 
 void ServiceRegistry::RemoveServiceConnector(
     ServiceConnectorBase* service_connector) {
+  RemoveServiceConnectorInternal(service_connector);
+  if (name_to_service_connector_.empty())
+    remote_service_provider_.reset();
+}
+
+bool ServiceRegistry::RemoveServiceConnectorInternal(
+    ServiceConnectorBase* service_connector) {
   NameToServiceConnectorMap::iterator it =
       name_to_service_connector_.find(service_connector->name());
-  assert(it != name_to_service_connector_.end());
+  if (it == name_to_service_connector_.end())
+    return false;
   delete it->second;
   name_to_service_connector_.erase(it);
-  if (name_to_service_connector_.empty())
-    remote_service_provider_.reset();
+  return true;
 }
 
 void ServiceRegistry::BindRemoteServiceProvider(
diff --git a/mojo/public/cpp/application/lib/service_registry.h b/mojo/public/cpp/application/lib/service_registry.h
index 478cd05..a2f23cd 100644
--- a/mojo/public/cpp/application/lib/service_registry.h
+++ b/mojo/public/cpp/application/lib/service_registry.h
@@ -5,6 +5,9 @@
 #ifndef MOJO_PUBLIC_CPP_APPLICATION_LIB_SERVICE_REGISTRY_H_
 #define MOJO_PUBLIC_CPP_APPLICATION_LIB_SERVICE_REGISTRY_H_
 
+#include <map>
+#include <string>
+
 #include "mojo/public/interfaces/service_provider/service_provider.mojom.h"
 
 namespace mojo {
@@ -40,6 +43,9 @@ class ServiceRegistry : public ServiceProvider {
       MOJO_OVERRIDE;
 
  private:
+  bool RemoveServiceConnectorInternal(
+      ServiceConnectorBase* service_connector);
+
   Application* application_;
   typedef std::map<std::string, ServiceConnectorBase*>
       NameToServiceConnectorMap;
diff --git a/mojo/public/cpp/application/tests/service_registry_unittest.cc b/mojo/public/cpp/application/tests/service_registry_unittest.cc
new file mode 100644
index 0000000..586dff8
--- /dev/null
+++ b/mojo/public/cpp/application/tests/service_registry_unittest.cc
@@ -0,0 +1,66 @@
+// Copyright 2014 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 "mojo/public/cpp/application/lib/service_registry.h"
+
+#include "mojo/public/cpp/application/lib/service_connector.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace mojo {
+namespace internal {
+namespace {
+
+class TestConnector : public ServiceConnectorBase {
+ public:
+  TestConnector(const std::string& name, int* delete_count)
+      : ServiceConnectorBase(name), delete_count_(delete_count) {}
+  virtual ~TestConnector() { (*delete_count_)++; }
+  virtual void ConnectToService(
+      const std::string& url,
+      const std::string& name,
+      ScopedMessagePipeHandle client_handle) MOJO_OVERRIDE {}
+ private:
+  int* delete_count_;
+};
+
+TEST(ServiceRegistryTest, Ownership) {
+  int delete_count = 0;
+
+  // Destruction.
+  {
+    ServiceRegistry registry(NULL);
+    registry.AddServiceConnector(new TestConnector("TC1", &delete_count));
+  }
+  EXPECT_EQ(1, delete_count);
+
+  // Removal.
+  {
+    ServiceRegistry registry(NULL);
+    ServiceConnectorBase* c = new TestConnector("TC1", &delete_count);
+    registry.AddServiceConnector(c);
+    registry.RemoveServiceConnector(c);
+    EXPECT_EQ(2, delete_count);
+  }
+
+  // Multiple.
+  {
+    ServiceRegistry registry(NULL);
+    registry.AddServiceConnector(new TestConnector("TC1", &delete_count));
+    registry.AddServiceConnector(new TestConnector("TC2", &delete_count));
+  }
+  EXPECT_EQ(4, delete_count);
+
+  // Re-addition.
+  {
+    ServiceRegistry registry(NULL);
+    registry.AddServiceConnector(new TestConnector("TC1", &delete_count));
+    registry.AddServiceConnector(new TestConnector("TC1", &delete_count));
+    EXPECT_EQ(5, delete_count);
+  }
+  EXPECT_EQ(6, delete_count);
+}
+
+}  // namespace
+}  // namespace internal
+}  // namespace mojo