Scan YAML framing preambles once

show_test.go

@@ -20,7 +20,7 @@ func TestEffectiveYAMLFramingPreservesOneNativeDocument(t *testing.T) {
 			if err := yaml.Unmarshal([]byte(native), &want); err != nil {
 				t.Fatal(err)
 			}
-			framed := frameEffectiveYAML("app", []byte(native))
+			framed, _ := frameEffectiveYAML("app", []byte(native))
 			withoutLabel := bytes.Replace(framed, []byte("# scope: app\n"), nil, 1)
 			if !strings.HasPrefix(native, "---") && !strings.HasPrefix(native, "%") {
 				withoutLabel = bytes.TrimPrefix(withoutLabel, []byte("---\n"))
@@ -45,3 +45,50 @@ func TestEffectiveYAMLFramingPreservesOneNativeDocument(t *testing.T) {
 		})
 	}
 }
+
+func TestEffectiveYAMLFramingPreservesBoundariesAndDirectives(t *testing.T) {
+	for _, test := range []struct {
+		name      string
+		scope     string
+		native    string
+		directive bool
+		label     string
+	}{
+		{name: "no trailing newline", scope: "app", native: "value: plain", label: "# scope: app\n"},
+		{name: "CRLF", scope: "app", native: "# authored\r\n---\r\nvalue: yes\r\n", label: "# scope: app\n"},
+		{name: "directive preamble", scope: "app", native: "%YAML 1.1\n# authored\n---\nvalue: yes", directive: true, label: "# scope: app\n"},
+		{name: "control character label", scope: "line\nbreak", native: "value: yes", label: "# scope: \"line\\nbreak\"\n"},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			var want any
+			if err := yaml.Unmarshal([]byte(test.native), &want); err != nil {
+				t.Fatal(err)
+			}
+			framed, hasDirective := frameEffectiveYAML(test.scope, []byte(test.native))
+			if hasDirective != test.directive {
+				t.Fatalf("hasDirective = %t, want %t", hasDirective, test.directive)
+			}
+			withoutLabel := bytes.Replace(framed, []byte(test.label), nil, 1)
+			if !strings.HasPrefix(test.native, "---") && !strings.HasPrefix(test.native, "%") {
+				withoutLabel = bytes.TrimPrefix(withoutLabel, []byte("---\n"))
+			}
+			if !bytes.Equal(withoutLabel, []byte(test.native)) {
+				t.Fatalf("native content changed: got %q want %q", withoutLabel, test.native)
+			}
+			if !bytes.Contains(framed, []byte(test.label)) {
+				t.Fatalf("missing scope label %q in %q", test.label, framed)
+			}
+			decoder := yaml.NewDecoder(bytes.NewReader(framed))
+			var first any
+			if err := decoder.Decode(&first); err != nil {
+				t.Fatal(err)
+			}
+			if !reflect.DeepEqual(first, want) {
+				t.Fatalf("decoded value = %#v, want %#v", first, want)
+			}
+			if err := decoder.Decode(new(any)); err != io.EOF {
+				t.Fatalf("unexpected extra YAML document: %v", err)
+			}
+		})
+	}
+}