#include "tensorflow/c/eager/c_api_test_util.h"
#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
#include "tensorflow/core/common_runtime/eager/eager_operation.h"
#include "tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h"
#include "tensorflow/core/platform/casts.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/cluster.pb.h"

namespace gcs = google::cloud::storage;

static std::string InitializeTmpDir() {
  // This env should be something like `gs://bucket/path`
  const char* test_dir = getenv("GCS_TEST_TMPDIR");
  if (test_dir != nullptr) {
    std::string bucket, object;
    TF_Status* status = TF_NewStatus();
    ParseGCSPath(test_dir, true, &bucket, &object, status);
    if (TF_GetCode(status) != TF_OK) {
      TF_DeleteStatus(status);
      return "";
    }

  ExpectScalarEq<float>(result_components[1].get(), 3.);
}

TEST(PARALLEL_DEVICE, TestCollectiveSync) { TestCollective(/*async=*/false); }

// Note that ops on the parallel device currently don't execute
// asynchronously. The test is just that we don't get deadlocks.
TEST(PARALLEL_DEVICE, TestCollectiveAsync) { TestCollective(/*async=*/true); }

// After the graph pass is registered, it takes effect globally and can affect
// other test cases. Define a static variable to switch it on and off.
bool GraphErrorInjectionPass::enabled_ = false;

// Test to ensure that a registered graph optimization pass is only executed
// once (i.e., on the main function side) in running distributed functions.
// This test creates a cluster with two workers, create a variable on the