"      value: 'add2:z:0'"
      "    }",
      &def));
  return def.SerializeAsString();
}

void TestFunctionWithPackedInput(const bool remote) {
  tensorflow::ServerDef server_def = GetServerDef(3);

  // This server def has the task index set to 0.
  string serialized = server_def.SerializeAsString();

  server_def.set_task_index(1);
  std::unique_ptr<tensorflow::GrpcServer> worker_server1;

BAZEL_COMMAND_RE = re.compile(
    r'(^| )(?P<command>bazel (.*? )?(?P<type>test|build) .+)')


class InvokeStatus:
  tests_failed = 'tests_failed'
  build_failed = 'build_failed'
  passed = 'passed'


def parse_args() -> argparse.Namespace:
  """Parses the commandline args."""
  parser = argparse.ArgumentParser(
      description='Extracts ResultStore links from a build log.\n'

// message) to `output_graph_def` (allocated by TF_NewBuffer()).
// `output_graph_def`'s underlying buffer will be freed when TF_DeleteBuffer()
// is called.
//
// May fail on very large graphs in the future.
TF_CAPI_EXPORT extern void TF_GraphToGraphDef(TF_Graph* graph,
                                              TF_Buffer* output_graph_def,
                                              TF_Status* status);

                       op_buf->length);
  ASSERT_EQ(TF_OK, TF_GetCode(status));

  const OpDef* expected_op_def;
  TF_ASSERT_OK(OpRegistry::Global()->LookUpOpDef("Add", &expected_op_def));
  string expected_serialized;
  expected_op_def->SerializeToString(&expected_serialized);
  EXPECT_EQ(expected_serialized, actual_string);
  TF_DeleteBuffer(g_buf);
  TF_DeleteBuffer(op_buf);

TFE_Context* CreateContext(const std::string& serialized_server_def,
                           bool isolate_session_state,
                           int64_t init_timeout_in_ms);

tensorflow::ServerDef ReplaceTaskInServerDef(
    const tensorflow::ServerDef& server_def, int task_index);

void ReplaceTaskInServerDef(tensorflow::ServerDef* server_def, int task_index,
                            const std::string& host, int port);

  // available to be called/referenced by its name in this context.
  virtual Status RegisterFunction(AbstractFunction*) = 0;
  // Remove a function. 'func' argument is the name of a previously added
  // FunctionDef. The name is in fdef.signature.name.
  virtual Status RemoveFunction(const string& func) = 0;

 private:
  const AbstractContextKind kind_;
};

namespace internal {
struct AbstractContextDeleter {