namespace mhlo {
namespace test {

// Given a raw string, return a ModuleOp that can be used with the given
// MLIRContext.
absl::StatusOr<OwningOpRef<ModuleOp>> GetMlirModuleFromString(
    absl::string_view module_string, MLIRContext* mlir_context);

}  // namespace test
}  // namespace mhlo
}  // namespace mlir

    return
  }
})";

class LegalizationOpConfigTest : public ::testing::Test {
 public:
  absl::Status CreateMlirModule(std::string module_string = kMlirModuleStr) {
    TF_ASSIGN_OR_RETURN(
        module_, test::GetMlirModuleFromString(module_string, &context_));

    context_.loadAllAvailableDialects();
    return absl::OkStatus();
  }

  absl::StatusOr<FuncOp> GetMain() {