public:
  TensorflowDialectToExecutorTest() {
    mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();
  }

  absl::Status CreateMlirModule(std::string mlir_module_filename) {
    std::string mlir_module_path = TestDataPath() + mlir_module_filename;
    mlir_module_ =
        mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);

};

TEST_F(LowerClusterToRuntimeOpsTest, SanityCheck) {
  TF_ASSERT_OK(CreateMlirModule("empty_func.mlir"));

  TF_EXPECT_OK(RunLowerClusterToRuntimeOpsPassPipeline(
      *mlir_module_, DeviceType(DEVICE_TPU_XLA_JIT)));
}

TEST_F(LowerClusterToRuntimeOpsTest, LowersClusterOpsTPU) {
  TF_ASSERT_OK(CreateMlirModule("basic_cluster.mlir"));

  TF_EXPECT_OK(RunLowerClusterToRuntimeOpsPassPipeline(

 public:
  SessionClusterTensorflowDialectTest() {
    mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();
  }

  absl::Status CreateMlirModule(std::string mlir_module_filename) {
    std::string mlir_module_path = TestDataPath() + mlir_module_filename;
    mlir_module_ =
        mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);

  TF_ASSERT_OK(CreateMlirModule("empty_func.mlir"));

  TF_EXPECT_OK(ExportFromTensorflowDialectToExecutor(*mlir_module_));

  EXPECT_EQ(compilation_status.Delta("success"), 1);
}

TEST_F(TensorflowDialectToExecutorTest, ErrorsWhenCannotConvert) {
  CellReader<int64_t> compilation_status(kCompilationStreamz);

  TF_ASSERT_OK(CreateMlirModule("invalid_executor.mlir"));

 public:
  FunctionClusterTensorflowDialectTest() {
    mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();
  }

  absl::Status CreateMlirModule(std::string mlir_module_filename) {
    std::string mlir_module_path = TestDataPath() + mlir_module_filename;
    mlir_module_ =
        mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);

// FILECHECK and MLIR tests.
class Tf2XlaRewriterTest : public ::testing::Test {
 public:
  void SetUp() override {
    tensorflow::XlaOpRegistry::RegisterCompilationKernels();
  }

  Status CreateMlirModule(std::string module_string = kMlirModuleStr) {
    TF_ASSIGN_OR_RETURN(
        module_, test::GetMlirModuleFromString(module_string, &context_));

    context_.loadAllAvailableDialects();

    env_ = Env::Default();

    test_group_name_ = "TestGroup";
    test_dir_ = testing::TmpDir();
    setenv("TF_DUMP_GRAPH_PREFIX", test_dir_.c_str(), 1);
  }

  absl::Status CreateMlirModule(std::string mlir_module_filename) {
    std::string mlir_module_path = TestDataPath() + mlir_module_filename;
    mlir_module_ =
        mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);

      absl::string_view mlir_module_filename) {
    TF_RETURN_IF_ERROR(CreateMlirModule(mlir_module_filename));
    std::vector<mlir::tf_device::ClusterFuncOp> cluster_func_ops;

    mlir_module_->walk([&](mlir::tf_device::ClusterFuncOp op) {
      cluster_func_ops.push_back(op);
    });
    return cluster_func_ops;
  }

 private:
  absl::Status CreateMlirModule(absl::string_view mlir_module_filename) {

  OwningOpRef<ModuleOp> module_;
};

TEST_F(LegalizationOpConfigTest, FailsWithExpectsLegalizationWithMlir) {
  TF_EXPECT_OK(CreateMlirModule());
  EXPECT_FALSE(IsOpLegalizedWithMlir(*module_->getOperation()));
}

TEST_F(LegalizationOpConfigTest, ExpectsFalseForNonMlirOps) {
  TF_EXPECT_OK(CreateMlirModule());
  TF_ASSERT_OK_AND_ASSIGN(FuncOp main, GetMain());

  main.walk([&](Operation* op) { EXPECT_FALSE(IsOpLegalizedWithMlir(*op)); });