TF_GraphCopyFunction(host_graph_, func_, nullptr, s_);
  ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);

  // Test we can retrieve function OpDef from graph
  TF_Buffer* buffer = TF_NewBuffer();
  TF_GraphGetOpDef(host_graph_, func_name_, buffer, s_);
  ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);

  // Sanity check returned OpDef
  std::string data(static_cast<const char*>(buffer->data), buffer->length);

  TF_DeleteStatus(status);
}

TEST(CAPI, GetOpDef) {
  TF_Status* status = TF_NewStatus();
  TF_Graph* graph = TF_NewGraph();
  TF_Buffer* buffer = TF_NewBuffer();

  TF_GraphGetOpDef(graph, "Add", buffer, status);
  ASSERT_EQ(TF_OK, TF_GetCode(status));
  const OpDef* expected_op_def;
  TF_ASSERT_OK(OpRegistry::Global()->LookUpOpDef("Add", &expected_op_def));
  std::string expected_serialized;

// Returns the serialized OpDef proto with name `op_name`, or a bad status if no
// such op exists. This can return OpDefs of functions copied into the graph.
TF_CAPI_EXPORT extern void TF_GraphGetOpDef(TF_Graph* graph,
                                            const char* op_name,
                                            TF_Buffer* output_op_def,
                                            TF_Status* status);

  GraphDef def;
  {
    mutex_lock l(graph->mu);
    graph->graph.ToGraphDef(&def);
  }
  status->status = MessageToBuffer(def, output_graph_def);
}

void TF_GraphGetOpDef(TF_Graph* graph, const char* op_name,
                      TF_Buffer* output_op_def, TF_Status* status) {
  const OpDef* op_def;
  {
    mutex_lock l(graph->mu);