diff --git a/buffet/commands/dbus_command_dispatcher_unittest.cc b/buffet/commands/dbus_command_dispatcher_unittest.cc
index 9808363..314135f 100644
--- a/buffet/commands/dbus_command_dispatcher_unittest.cc
+++ b/buffet/commands/dbus_command_dispatcher_unittest.cc
@@ -108,9 +108,8 @@
   }
 
   DBusCommandProxy* FindProxy(CommandInstance* command_instance) {
-    const auto& command_map = command_dispatcher_->command_map_;
-    auto it = command_map.find(command_instance);
-    return it != command_map.end() ? it->second.get() : nullptr;
+    CHECK_EQ(command_instance->proxies_.size(), 1);
+    return static_cast<DBusCommandProxy*>(command_instance->proxies_[0].get());
   }
 
   void FinishCommand(DBusCommandProxy* proxy) {
@@ -156,7 +155,6 @@
   EXPECT_CALL(*mock_exported_object_manager_, SendSignal(_)).Times(2);
   FinishCommand(command_proxy);
 
-  EXPECT_EQ(nullptr, FindProxy(command_instance));
   EXPECT_EQ(nullptr, command_queue_.Find(id));
 }
 
@@ -190,7 +188,6 @@
   EXPECT_CALL(*mock_exported_object_manager_, SendSignal(_)).Times(2);
   FinishCommand(command_proxy);
 
-  EXPECT_EQ(nullptr, FindProxy(command_instance));
   EXPECT_EQ(nullptr, command_queue_.Find(id));
 }