rewrite_fn(arg_source_tensors, &g, nullptr, nullptr, &call_node_def));
  node_name_image = g->BuildNodeNameIndex();

  // Verify step 1: add key placeholder node.
  Node *key_placeholder = node_name_image["cluster_key_placeholder"];
  EXPECT_NE(key_placeholder, nullptr);
  // Verify step 2: replace _Arg nodes with XlaRecvAtHost.
  for (Node *n : g->nodes()) {
    EXPECT_NE(n->type_string(), "_Arg");
  }