#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(exported_model.graph_def().node(), SizeIs(1));
  EXPECT_THAT(exported_model.graph_def().node()[0].name(), StrEq("foo"));

  EXPECT_THAT(exported_model.init_node_name(), StrEq("init_node_name"));
  EXPECT_THAT(exported_model.checkpoint_dir(), StrEq("checkpoint_dir"));
  EXPECT_FALSE(exported_model.has_saver_def());
  EXPECT_THAT(exported_model.function_aliases(), IsEmpty());

  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) {

  EXPECT_THAT(default_spec.matcher().function_name().regex(), StrEq(".*"));
  EXPECT_TRUE(default_spec.method().has_static_range_ptq());

  // Test that the expansion for convolution ops is done.
  const QuantizationSpec& conv_spec = new_config.specs().specs(1);
  EXPECT_THAT(conv_spec.matcher().function_name().regex(),
              StrEq("composite_conv.*"));
  ASSERT_TRUE(conv_spec.method().has_static_range_ptq());

						s = field[1];
						if (!s || streq(s, "POSIX"))
							s = "C";
						if ((ans = setlocale(LC_COLLATE, s)) && streq(ans, "POSIX"))
							ans = "C";
						if (!ans || !streq(ans, s) && streq(s, "C"))
							ans = 0;
						else if ((ans = setlocale(LC_CTYPE, s)) && streq(ans, "POSIX"))
							ans = "C";
						if (!ans || !streq(ans, s) && streq(s, "C"))
							skip = note(level, s, skip, test);

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);

              StrEq("composite_dot_general_fn_2"));
  EXPECT_TRUE(quantized_result.method().has_static_range_ptq());

  // Test that the non-quantized op is captured in `results`.
  const QuantizationResult& non_quantized_result = results.results(1);
  EXPECT_THAT(non_quantized_result.quantizable_unit().name(),
              StrEq("composite_dot_general_fn_1"));

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;

  SmallVector<func::FuncOp> funcs = GetSortedFunctions(*module_op);
  ASSERT_THAT(funcs, SizeIs(3));
  EXPECT_THAT(funcs[0].getSymName(), StrEq("conv_1_fn"));
  EXPECT_THAT(funcs[1].getSymName(), StrEq("conv_2_fn"));
  EXPECT_THAT(funcs[2].getSymName(), StrEq("conv_3_fn"));
}

}  // namespace