From a6fbca2eec789ed5ca0704eb0816021ead957902 Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Mon, 23 May 2022 19:09:10 -0700
Subject: [PATCH] `yaml_test`: Add argument to allow non-strict YAML decoding.

PiperOrigin-RevId: 450576605
---
 tools/yamltest/defs.bzl | 8 +++++++-
 tools/yamltest/main.go  | 9 ++++++---
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/tools/yamltest/defs.bzl b/tools/yamltest/defs.bzl
index fd04f947d..04469d815 100644
--- a/tools/yamltest/defs.bzl
+++ b/tools/yamltest/defs.bzl
@@ -6,9 +6,10 @@ def _yaml_test_impl(ctx):
     ctx.actions.write(runner, "\n".join([
         "#!/bin/bash",
         "set -euo pipefail",
-        "%s -schema=%s -- %s" % (
+        "%s -schema=%s -strict=%s -- %s" % (
             ctx.files._tool[0].short_path,
             ctx.files.schema[0].short_path,
+            "true" if ctx.attr.strict else "false",
             " ".join([f.short_path for f in ctx.files.srcs]),
         ),
     ]), is_executable = True)
@@ -31,6 +32,11 @@ yaml_test = rule(
             allow_single_file = True,
             mandatory = True,
         ),
+        "strict": attr.bool(
+            doc = "Whether to use strict mode for YAML decoding.",
+            mandatory = False,
+            default = True,
+        ),
         "_tool": attr.label(
             executable = True,
             cfg = "host",
diff --git a/tools/yamltest/main.go b/tools/yamltest/main.go
index 88271fb66..b9d8cc3db 100644
--- a/tools/yamltest/main.go
+++ b/tools/yamltest/main.go
@@ -26,6 +26,11 @@ import (
 	yaml "gopkg.in/yaml.v2"
 )
 
+var (
+	schema = flag.String("schema", "", "path to JSON schema file.")
+	strict = flag.Bool("strict", true, "Whether to enable strict mode for YAML decoding")
+)
+
 func fixup(v interface{}) (interface{}, error) {
 	switch x := v.(type) {
 	case map[interface{}]interface{}:
@@ -65,7 +70,7 @@ func loadFile(filename string) (gojsonschema.JSONLoader, error) {
 	}
 	defer f.Close()
 	dec := yaml.NewDecoder(f)
-	dec.SetStrict(true)
+	dec.SetStrict(*strict)
 	var object interface{}
 	if err := dec.Decode(&object); err != nil {
 		return nil, err
@@ -81,8 +86,6 @@ func loadFile(filename string) (gojsonschema.JSONLoader, error) {
 	return gojsonschema.NewStringLoader(string(bytes)), nil
 }
 
-var schema = flag.String("schema", "", "path to JSON schema file.")
-
 func main() {
 	flag.Parse()
 	if *schema == "" || len(flag.Args()) == 0 {