def test_write_load_stats_text(self): stats = text_format.Parse(""" datasets {} """, statistics_pb2.DatasetFeatureStatisticsList()) stats_path = os.path.join(FLAGS.test_tmpdir, 'stats.pbtxt') stats_util.write_stats_text(stats=stats, output_path=stats_path) loaded_stats = stats_util.load_stats_text(input_path=stats_path) self.assertEqual(stats, loaded_stats)
def test_write_stats_text_invalid_stats_input(self): with self.assertRaisesRegex( TypeError, '.*should be a DatasetFeatureStatisticsList proto.'): _ = stats_util.write_stats_text({}, 'stats.pbtxt')