constexpr char kFailure[] = "kFailure";

class MockMlirOptimizationPass : public MlirOptimizationPass {
 public:
  MOCK_METHOD(llvm::StringRef, name, (), (const, override));
  MOCK_METHOD(MlirOptimizationPassState, GetPassState,
              (const DeviceSet* device_set, const ConfigProto& config_proto,
               const Graph& graph,
               const FunctionLibraryDefinition& function_library),
              (const, override));

                                "The number of graphs being op expanded.");
}  // namespace

namespace tfr {

MlirOptimizationPassState GraphDecomposePass::GetPassState(
    const DeviceSet* device_set, const ConfigProto& config_proto,
    const Graph& graph,
    const FunctionLibraryDefinition& function_library) const {
  const char* tfr_lib_env_val = getenv(std::string(kTFRLibEnv).c_str());

class MlirGraphOptimizationPass : public ::tensorflow::MlirOptimizationPass {
 public:
  llvm::StringRef name() const override { return "graph_optimization"; }

  ::tensorflow::MlirOptimizationPassState GetPassState(
      const ::tensorflow::DeviceSet* device_set,
      const ::tensorflow::ConfigProto& config_proto,
      const tensorflow::Graph& graph,
      const tensorflow::FunctionLibraryDefinition& function_library)

  // `function_library` contains function definitions for function calls in
  // `graph` not included in the `graph` FunctionLibraryDefinition.
  virtual MlirOptimizationPassState GetPassState(
      const DeviceSet* device_set, const ConfigProto& config_proto,
      const Graph& graph,
      const FunctionLibraryDefinition& function_library) const = 0;

  virtual Status Run(const std::string& function_name,

  // Whether to run this pass. If this is enabled, the GraphDef will be imported
  // to MLIR even no tf composition file is found.
  ::tensorflow::MlirOptimizationPassState GetPassState(
      const DeviceSet* device_set, const ConfigProto& config_proto,
      const Graph& graph,
      const FunctionLibraryDefinition& function_library) const override;

Status MlirGraphOptimizationPass::Run(
    const std::string& function_name, const ConfigProto& config_proto,
    ModuleOp module, const Graph& graph,
    const tensorflow::FunctionLibraryDefinition& function_library) {
  if (GetPassState(/*device_set=*/nullptr, config_proto, graph,
                   function_library) ==
      ::tensorflow::MlirOptimizationPassState::Disabled) {
    VLOG(1) << "Skipping MLIR Graph Optimization Pass"

  int num_passes_enabled = 0, num_passes_disabled = 0,
      num_passes_fallback_enabled = 0;
  for (const auto& pass_registration : registry_->passes()) {
    MlirOptimizationPassState pass_state = pass_registration.pass->GetPassState(
        &device_set, config_proto, **graph, *flib_def);
    per_pass_state.push_back(pass_state);
    switch (pass_state) {
      case MlirOptimizationPassState::FallbackEnabled: {