mlir::OwningOpRef<mlir::ModuleOp> module =
      mlir::parseSourceString<mlir::ModuleOp>(code, &context);
  ASSERT_TRUE(module);
  EXPECT_TRUE(HasTPUPartitionedCallOpInModule(*module));
}

TEST(HasTPUPartitionedCallOpInModule, HasNotTPUPartitionedCallModule) {
  const char* const code = R"mlir(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {

  return IsReplicatedGraph(graph, function_library);
}

bool IsSupportedByReplicatedBridge(mlir::ModuleOp module) {
  return IsReplicatedGraph(module);
}

bool HasTPUPartitionedCallOpInModule(mlir::ModuleOp module) {
  bool has_tpu_partitioned_call = false;
  for (auto func_op : module.getOps<mlir::func::FuncOp>()) {
    func_op->walk([&](mlir::TF::TPUPartitionedCallOp op) {