diff --git a/src/component_manager_unittest.cc b/src/component_manager_unittest.cc
index 83f9933..31949d7 100644
--- a/src/component_manager_unittest.cc
+++ b/src/component_manager_unittest.cc
@@ -31,38 +31,38 @@
 }
 
 // Creates sample trait/component trees:
-// {
-//   "traits": {
-//     "t1": {},
-//     "t2": {},
-//     "t3": {},
-//     "t4": {},
-//     "t5": {},
-//     "t6": {},
-//   },
-//   "components": {
-//     "comp1": {
-//       "traits": [ "t1" ],
-//       "components": {
-//         "comp2": [
-//           { "traits": [ "t2" ] },
-//           {
-//             "traits": [ "t3" ],
-//             "components": {
-//               "comp3": {
-//                 "traits": [ "t4" ],
-//                 "components": {
-//                   "comp4": {
-//                     "traits": [ "t5", "t6" ]
-//                   }
-//                 }
-//               }
-//             }
-//           }
-//         ],
-//       }
-//     }
-//   }
+// {
+//   "traits": {
+//     "t1": {},
+//     "t2": {},
+//     "t3": {},
+//     "t4": {},
+//     "t5": {},
+//     "t6": {},
+//   },
+//   "components": {
+//     "comp1": {
+//       "traits": [ "t1" ],
+//       "components": {
+//         "comp2": [
+//           { "traits": [ "t2" ] },
+//           {
+//             "traits": [ "t3" ],
+//             "components": {
+//               "comp3": {
+//                 "traits": [ "t4" ],
+//                 "components": {
+//                   "comp4": {
+//                     "traits": [ "t5", "t6" ]
+//                   }
+//                 }
+//               }
+//             }
+//           }
+//         ],
+//       }
+//     }
+//   }
 // }
 void CreateTestComponentTree(ComponentManager* manager) {
   const char kTraits[] = R"({"t1":{},"t2":{},"t3":{},"t4":{},"t5":{},"t6":{}})";
@@ -670,7 +670,7 @@
   // Component comp1 doesn't have trait2.
   EXPECT_FALSE(manager.AddCommand(*command3, UserRole::kOwner, &id, nullptr));
 
-  // No component specified, use the first top-level component (comp1)
+  // No component specified, find the suitable component
   const char kCommand4[] = R"({
     "name": "trait1.command1",
     "parameters": {}
@@ -680,6 +680,16 @@
   auto cmd = manager.FindCommand(id);
   ASSERT_NE(nullptr, cmd);
   EXPECT_EQ("comp1", cmd->GetComponent());
+
+  const char kCommand5[] = R"({
+    "name": "trait2.command1",
+    "parameters": {}
+  })";
+  auto command5 = CreateDictionaryValue(kCommand5);
+  EXPECT_TRUE(manager.AddCommand(*command5, UserRole::kOwner, &id, nullptr));
+  cmd = manager.FindCommand(id);
+  ASSERT_NE(nullptr, cmd);
+  EXPECT_EQ("comp2", cmd->GetComponent());
 }
 
 TEST(ComponentManager, AddCommandHandler) {
@@ -750,30 +760,30 @@
   auto state1 = CreateDictionaryValue(kState1);
   ASSERT_TRUE(manager.SetStateProperties("comp1", *state1, nullptr));
   const char kExpected1[] = R"({
-    "comp1": {
-      "traits": [ "t1" ],
-      "state": {"t1": {"p1": 0, "p2": "foo"}},
-      "components": {
-        "comp2": [
-          {
-            "traits": [ "t2" ]
-          },
-          {
-            "traits": [ "t3" ],
-            "components": {
-              "comp3": {
-                "traits": [ "t4" ],
-                "components": {
-                  "comp4": {
-                    "traits": [ "t5", "t6" ]
-                  }
-                }
-              }
-            }
-          }
-        ]
-      }
-    }
+    "comp1": {
+      "traits": [ "t1" ],
+      "state": {"t1": {"p1": 0, "p2": "foo"}},
+      "components": {
+        "comp2": [
+          {
+            "traits": [ "t2" ]
+          },
+          {
+            "traits": [ "t3" ],
+            "components": {
+              "comp3": {
+                "traits": [ "t4" ],
+                "components": {
+                  "comp4": {
+                    "traits": [ "t5", "t6" ]
+                  }
+                }
+              }
+            }
+          }
+        ]
+      }
+    }
   })";
   EXPECT_JSON_EQ(kExpected1, manager.GetComponents());
 
@@ -782,30 +792,30 @@
   ASSERT_TRUE(manager.SetStateProperties("comp1", *state2, nullptr));
 
   const char kExpected2[] = R"({
-    "comp1": {
-      "traits": [ "t1" ],
-      "state": {"t1": {"p1": {"bar": "baz"}, "p2": "foo"}},
-      "components": {
-        "comp2": [
-          {
-            "traits": [ "t2" ]
-          },
-          {
-            "traits": [ "t3" ],
-            "components": {
-              "comp3": {
-                "traits": [ "t4" ],
-                "components": {
-                  "comp4": {
-                    "traits": [ "t5", "t6" ]
-                  }
-                }
-              }
-            }
-          }
-        ]
-      }
-    }
+    "comp1": {
+      "traits": [ "t1" ],
+      "state": {"t1": {"p1": {"bar": "baz"}, "p2": "foo"}},
+      "components": {
+        "comp2": [
+          {
+            "traits": [ "t2" ]
+          },
+          {
+            "traits": [ "t3" ],
+            "components": {
+              "comp3": {
+                "traits": [ "t4" ],
+                "components": {
+                  "comp4": {
+                    "traits": [ "t5", "t6" ]
+                  }
+                }
+              }
+            }
+          }
+        ]
+      }
+    }
   })";
   EXPECT_JSON_EQ(kExpected2, manager.GetComponents());
 
@@ -815,31 +825,31 @@
                                          nullptr));
 
   const char kExpected3[] = R"({
-    "comp1": {
-      "traits": [ "t1" ],
-      "state": {"t1": {"p1": {"bar": "baz"}, "p2": "foo"}},
-      "components": {
-        "comp2": [
-          {
-            "traits": [ "t2" ]
-          },
-          {
-            "traits": [ "t3" ],
-            "components": {
-              "comp3": {
-                "traits": [ "t4" ],
-                "components": {
-                  "comp4": {
-                    "traits": [ "t5", "t6" ],
-                    "state": { "t5": { "p1": 1 } }
-                  }
-                }
-              }
-            }
-          }
-        ]
-      }
-    }
+    "comp1": {
+      "traits": [ "t1" ],
+      "state": {"t1": {"p1": {"bar": "baz"}, "p2": "foo"}},
+      "components": {
+        "comp2": [
+          {
+            "traits": [ "t2" ]
+          },
+          {
+            "traits": [ "t3" ],
+            "components": {
+              "comp3": {
+                "traits": [ "t4" ],
+                "components": {
+                  "comp4": {
+                    "traits": [ "t5", "t6" ],
+                    "state": { "t5": { "p1": 1 } }
+                  }
+                }
+              }
+            }
+          }
+        ]
+      }
+    }
   })";
   EXPECT_JSON_EQ(kExpected3, manager.GetComponents());
 }
@@ -853,33 +863,33 @@
       nullptr));
 
   const char kExpected[] = R"({
-    "comp1": {
-      "traits": [ "t1" ],
-      "components": {
-        "comp2": [
-          {
-            "traits": [ "t2" ]
-          },
-          {
-            "traits": [ "t3" ],
-            "components": {
-              "comp3": {
-                "traits": [ "t4" ],
-                "components": {
-                  "comp4": {
-                    "traits": [ "t5", "t6" ],
-                    "state": {
-                      "t5": { "p1": 3 },
-                      "t6": { "p2": 5 }
-                    }
-                  }
-                }
-              }
-            }
-          }
-        ]
-      }
-    }
+    "comp1": {
+      "traits": [ "t1" ],
+      "components": {
+        "comp2": [
+          {
+            "traits": [ "t2" ]
+          },
+          {
+            "traits": [ "t3" ],
+            "components": {
+              "comp3": {
+                "traits": [ "t4" ],
+                "components": {
+                  "comp4": {
+                    "traits": [ "t5", "t6" ],
+                    "state": {
+                      "t5": { "p1": 3 },
+                      "t6": { "p2": 5 }
+                    }
+                  }
+                }
+              }
+            }
+          }
+        ]
+      }
+    }
   })";
   EXPECT_JSON_EQ(kExpected, manager.GetComponents());
 }
@@ -908,12 +918,12 @@
   ASSERT_TRUE(manager.SetStateProperty("comp1", "trait1.prop1", p1, nullptr));
 
   const char kExpected1[] = R"({
-    "comp1": {
-      "traits": [ "trait1", "trait2" ],
-      "state": {
-        "trait1": { "prop1": "foo" }
-      }
-    }
+    "comp1": {
+      "traits": [ "trait1", "trait2" ],
+      "state": {
+        "trait1": { "prop1": "foo" }
+      }
+    }
   })";
   EXPECT_JSON_EQ(kExpected1, manager.GetComponents());
 
@@ -921,13 +931,13 @@
   ASSERT_TRUE(manager.SetStateProperty("comp1", "trait2.prop3", p2, nullptr));
 
   const char kExpected2[] = R"({
-    "comp1": {
-      "traits": [ "trait1", "trait2" ],
-      "state": {
-        "trait1": { "prop1": "foo" },
-        "trait2": { "prop3": 2 }
-      }
-    }
+    "comp1": {
+      "traits": [ "trait1", "trait2" ],
+      "state": {
+        "trait1": { "prop1": "foo" },
+        "trait2": { "prop3": 2 }
+      }
+    }
   })";
   EXPECT_JSON_EQ(kExpected2, manager.GetComponents());
   // Just the package name without property:
@@ -1085,4 +1095,22 @@
   EXPECT_EQ(snapshot.update_id, updates2.front());
 }
 
+TEST(ComponentManager, FindComponentWithTrait) {
+  ComponentManager manager;
+  const char kTraits[] = R"({
+    "trait1": {},
+    "trait2": {},
+    "trait3": {}
+  })";
+  auto traits = CreateDictionaryValue(kTraits);
+  ASSERT_TRUE(manager.LoadTraits(*traits, nullptr));
+  ASSERT_TRUE(manager.AddComponent("", "comp1", {"trait1", "trait2"}, nullptr));
+  ASSERT_TRUE(manager.AddComponent("", "comp2", {"trait3"}, nullptr));
+
+  EXPECT_EQ("comp1", manager.FindComponentWithTrait("trait1"));
+  EXPECT_EQ("comp1", manager.FindComponentWithTrait("trait2"));
+  EXPECT_EQ("comp2", manager.FindComponentWithTrait("trait3"));
+  EXPECT_EQ("", manager.FindComponentWithTrait("trait4"));
+}
+
 }  // namespace weave