#include "tsl/platform/protobuf.h"  // IWYU pragma: keep

namespace stablehlo::quantization {
namespace {

using ::tensorflow::GraphDef;
using ::tensorflow::NodeDef;
using ::testing::SizeIs;
using ::testing::StrEq;
using ::tsl::protobuf::TextFormat;

TEST(GraphDefTest, MutateNodeDefsMutatesTopLevelNodeDefs) {
  GraphDef graph_def;
  ASSERT_TRUE(TextFormat::ParseFromString(R"pb(

  ASSERT_THAT(*representative_dataset_file_map, SizeIs(1));
  EXPECT_THAT(*representative_dataset_file_map,
              Contains(Key(StrEq("test_signature_key"))));
  EXPECT_THAT(representative_dataset_file_map->at("test_signature_key")
                  .tfrecord_file_path(),
              StrEq("test_path"));
}

TEST(CreateRepresentativeDatasetFileMapTest,
     ConfigWithDuplicateSignatureKeyReturnsInvalidArgumentError) {

using ::stablehlo::quantization::PopulateDefaults;
using ::stablehlo::quantization::QuantizationConfig;
using ::testing::Contains;
using ::testing::SizeIs;
using ::testing::StartsWith;
using ::testing::StrEq;
using ::tsl::testing::IsOk;

// Matches an operation whose `getSymName` equals `name`.
MATCHER_P(HasSymName, name, "") {
  auto non_const_arg = const_cast<std::remove_const_t<decltype(arg)>>(arg);

namespace mlir::quant::stablehlo {
namespace {

using ::stablehlo::quantization::QuantizationResults;
using ::stablehlo::quantization::io::ReadFileToString;
using ::testing::SizeIs;
using ::testing::StrEq;
using ::tsl::protobuf::TextFormat;
using ::tsl::testing::IsOk;
using ::tsl::testing::StatusIs;

using SaveQuantizationReportInstrumentationTest = QuantizationTestBase;