mlir_module_ =
        mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
    if (!mlir_module_) {
      return absl::Status(
          absl::StatusCode::kNotFound,
          absl::StrCat("Could not find MLIR module at ", mlir_module_path));
    }
    return absl::OkStatus();
  }

  DialectRegistry registry_;
  MLIRContext context_;

}

TEST(TensorPjRtBufferUtilTest, GetPjRtCApiClientNotFound) {
  EXPECT_THAT(
      GetPjRtCApiClient(tensorflow::DeviceType(DEVICE_CPU)),
      StatusIs(error::NOT_FOUND,
               HasSubstr(absl::StrCat("PjRt client not found for device type ",
                                      DEVICE_CPU))));
}

TEST(TensorPjRtBufferUtilTest, GetPjRtCApiClientIncorrectType) {

  // The report file is not created because `QuantizeCompositeFunctionsPass` was
  // not run.
  EXPECT_THAT(ReadFileToString(report_file_path),
              StatusIs(absl::StatusCode::kNotFound));
}

TEST_F(SaveQuantizationReportInstrumentationTest,
       ReportNotSavedWhenReportFilePathIsNullopt) {
  constexpr absl::string_view kModuleWithCompositeDotGeneral = R"mlir(

      t.scalar<tensorflow::ResourceHandle>()().set_device("invalid_device");

      outputs->push_back(t);
    } else if (absl::StartsWith(output_name, "var")) {
      return Status(absl::StatusCode::kNotFound,
                    "Can't find variable " + output_name + " in session");
    } else {
      // Create a scalar float tensor.
      Tensor t = Tensor(tensorflow::DT_FLOAT, tensorflow::TensorShape({}));

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

namespace tensorflow {
namespace metrics {

const char kFingerprintFound[] = "FOUND";
const char kFingerprintNotFound[] = "NOT_FOUND";
const char kFingerprintError[] = "ERROR";

// Returns "/tensorflow/core/saved_model/write/count" cell. This metric
// has 1 field "write_version", which is equal to the