using ::mlir::MLIRContext;
using ::mlir::ModuleOp;
using ::mlir::OwningOpRef;
using ::mlir::mhlo::test::GetMlirModuleFromString;
using ::tensorflow::monitoring::testing::CellReader;

static constexpr char kHasTpuPartitionedCallStreamzName[] =
    "/tensorflow/core/tf2xla/internal/inference/tpu_partitioned_call";

class InferenceMetricsPassTest : public ::testing::Test {
 protected:
  void CreateModule(const char* module_string) {