fix wait_for_message for custom ros contexts

rclcpp/include/rclcpp/wait_for_message.hpp

@@ -56,7 +56,7 @@ bool wait_for_message(
       }
     });
 
-  rclcpp::WaitSet wait_set;
+  rclcpp::WaitSet wait_set({}, {}, {}, {}, {}, {}, context);
   wait_set.add_subscription(subscription);
   RCPPUTILS_SCOPE_EXIT(wait_set.remove_subscription(subscription); );
   wait_set.add_guard_condition(gc);

rclcpp/test/rclcpp/test_wait_for_message.cpp

@@ -134,3 +134,32 @@ TEST(TestUtilities, wait_for_last_message) {
 
   rclcpp::shutdown();
 }
+
+TEST(TestUtilities, wait_for_message_custom_context) {
+  auto context = std::make_shared<rclcpp::Context>();
+  context->init(0, nullptr);
+
+  auto node_opt = rclcpp::NodeOptions().context(context);
+  auto node = std::make_shared<rclcpp::Node>("wait_for_message_custom_context_node", node_opt);
+
+  using MsgT = test_msgs::msg::Strings;
+  auto pub = node->create_publisher<MsgT>("wait_for_message_topic", 10);
+
+  MsgT out;
+  auto received = false;
+  auto wait = std::async(
+    [&]() {
+      auto ret = rclcpp::wait_for_message(out, node, "wait_for_message_topic", 5s);
+      EXPECT_TRUE(ret);
+      received = true;
+    });
+
+  for (auto i = 0u; i < 10 && received == false; ++i) {
+    pub->publish(*get_messages_strings()[0]);
+    std::this_thread::sleep_for(1s);
+  }
+  ASSERT_TRUE(received);
+  EXPECT_EQ(out, *get_messages_strings()[0]);
+
+  context->shutdown("test complete");
+}