using tracing::TracingOperation;

class CppGradients
    : public ::testing::TestWithParam<std::tuple<const char*, bool, bool>> {
 protected:
  void SetUp() override {
    TF_StatusPtr status(TF_NewStatus());
    TF_SetTracingImplementation(std::get<0>(GetParam()), status.get());
    absl::Status s = StatusFromTF_Status(status.get());
    CHECK_EQ(errors::OK, s.code()) << s.message();
  }
};

#include "tensorflow/core/protobuf/tensorflow_server.pb.h"

using tensorflow::string;

namespace {

void BM_InitOp(::testing::benchmark::State& state) {
  TF_Status* status = TF_NewStatus();
  TFE_ContextOptions* opts = TFE_NewContextOptions();
  TFE_Context* ctx = TFE_NewContext(opts, status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TFE_DeleteContextOptions(opts);

  TF_DataType type;
  TF_Status* s = TF_NewStatus();
  TF_OpKernelConstruction_GetAttrType(ctx, "foobar", &type, s);
  TF_DeleteStatus(s);
  return NULL;
}

// A compute function. This will never actually get called in this test, it's
// just nice to know that it compiles.
void compute(void* kernel, TF_OpKernelContext* ctx) {
  TF_Tensor* input;
  TF_Status* s = TF_NewStatus();
  TF_GetInput(ctx, 0, &input, s);

  TF_Buffer* op_list_buf = TF_GetAllOpList();
  TF_Status* status = TF_NewStatus();
  auto* api_def_map = TF_NewApiDefMap(op_list_buf, status);
  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TF_DeleteStatus(status);

  std::string op_name = "TestCApi";
  status = TF_NewStatus();
  auto* api_def_buf =
      TF_ApiDefMapGet(api_def_map, op_name.c_str(), op_name.size(), status);

                       int* num_retvals) override {
    auto* tf_opdesc = op_.release();
    if (tf_opdesc == nullptr) {
      return errors::InvalidArgument("AbstractOp is incomplete.");
    }
    TF_Status* s = TF_NewStatus();
    auto* operation = TF_FinishOperation(tf_opdesc, s);
    TF_RETURN_IF_ERROR(StatusFromTF_Status(s));
    TF_DeleteStatus(s);
    *num_retvals = TF_OperationNumOutputs(operation);

  return node_def.op() == "Neg" && node_def.name() == "neg" &&
         node_def.input_size() == 1 && node_def.input(0) == input;
}

bool GetGraphDef(TF_Graph* graph, tensorflow::GraphDef* graph_def) {
  TF_Status* s = TF_NewStatus();
  TF_Buffer* buffer = TF_NewBuffer();
  TF_GraphToGraphDef(graph, buffer, s);
  bool ret = TF_GetCode(s) == TF_OK;
  EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);

template <class T, TF_DataType datatype>
TFE_TensorHandle* TestTensorHandleWithDims(TFE_Context* ctx, const T* data,
                                           const int64_t* dims, int num_dims) {
  TF_Status* status = TF_NewStatus();
  TF_Tensor* t = TFE_AllocateHostTensor(ctx, datatype, dims, num_dims, status);
  memcpy(TF_TensorData(t), data, TF_TensorByteSize(t));
  TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);

  ASSERT_TRUE(tensorflow::GrpcServer::Create(
                  server_def, tensorflow::Env::Default(), &worker_server2)
                  .ok());
  ASSERT_TRUE(worker_server2->Start().ok());

  TF_Status* status = TF_NewStatus();
  TFE_ContextOptions* opts = TFE_NewContextOptions();
  TFE_ContextOptionsSetAsync(opts, static_cast<unsigned char>(/*enable=*/true));
  TFE_ContextOptionsSetDevicePlacementPolicy(opts, TFE_DEVICE_PLACEMENT_SILENT);

    EXPECT_EQ(data_out[j], data_in[j]);
  }

  TFE_CallDLManagedTensorDeleter(dlm_out);
  TFE_DeleteTensorHandle(handle);
}

TEST(DLPack, HandleFromDLPackStrides) {
  TF_Status* status = TF_NewStatus();
  TFE_ContextOptions* opts = TFE_NewContextOptions();
  TFE_Context* ctx = TFE_NewContext(opts, status);
  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TFE_DeleteContextOptions(opts);

  using Base::pair;

  std::string ToString() const { return absl::StrCat(first, "->", second); }
};

class CApiFunctionTest : public ::testing::Test {
 protected:
  CApiFunctionTest()
      : s_(TF_NewStatus()),
        func_graph_(TF_NewGraph()),
        host_graph_(TF_NewGraph()),
        func_(nullptr) {}

  void SetUp() override {}

  ~CApiFunctionTest() override {
    TF_DeleteFunction(func_);