From 4f8e8f4cacf11e4d69d587e30d99ae7d88cc0878 Mon Sep 17 00:00:00 2001
From: Owen Williams <owen.williams@grafana.com>
Date: Wed, 4 Sep 2024 11:04:59 -0400
Subject: [PATCH] expfmt: Add WithEscapingScheme to help construct Formats
 (#688)

expfmt: Add WithEscapingScheme to help construct Formats

Signed-off-by: Owen Williams <owen.williams@grafana.com>

---------

Signed-off-by: Owen Williams <owen.williams@grafana.com>
---
 expfmt/expfmt.go      | 23 +++++++++++++++++++++
 expfmt/expfmt_test.go | 47 +++++++++++++++++++++++++++++++++++++------
 2 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/expfmt/expfmt.go b/expfmt/expfmt.go
index 14034a67..d942af8e 100644
--- a/expfmt/expfmt.go
+++ b/expfmt/expfmt.go
@@ -112,6 +112,29 @@ func NewOpenMetricsFormat(version string) (Format, error) {
 	return FmtUnknown, fmt.Errorf("unknown open metrics version string")
 }
 
+// WithEscapingScheme returns a copy of Format with the specified escaping
+// scheme appended to the end. If an escaping scheme already exists it is
+// removed.
+func (f Format) WithEscapingScheme(s model.EscapingScheme) Format {
+	var terms []string
+	for _, p := range strings.Split(string(f), ";") {
+		toks := strings.Split(p, "=")
+		if len(toks) != 2 {
+			trimmed := strings.TrimSpace(p)
+			if len(trimmed) > 0 {
+				terms = append(terms, trimmed)
+			}
+			continue
+		}
+		key := strings.TrimSpace(toks[0])
+		if key != model.EscapingKey {
+			terms = append(terms, strings.TrimSpace(p))
+		}
+	}
+	terms = append(terms, model.EscapingKey+"="+s.String())
+	return Format(strings.Join(terms, "; "))
+}
+
 // FormatType deduces an overall FormatType for the given format.
 func (f Format) FormatType() FormatType {
 	toks := strings.Split(string(f), ";")
diff --git a/expfmt/expfmt_test.go b/expfmt/expfmt_test.go
index 3a16f3b5..bc946ec1 100644
--- a/expfmt/expfmt_test.go
+++ b/expfmt/expfmt_test.go
@@ -17,6 +17,8 @@ import (
 	"testing"
 
 	"github.com/prometheus/common/model"
+
+	"github.com/stretchr/testify/require"
 )
 
 // Test Format to Escapting Scheme conversion
@@ -92,9 +94,7 @@ func TestToFormatType(t *testing.T) {
 		},
 	}
 	for _, test := range tests {
-		if test.format.FormatType() != test.expected {
-			t.Errorf("expected %v got %v", test.expected, test.format.FormatType())
-		}
+		require.Equal(t, test.expected, test.format.FormatType())
 	}
 }
 
@@ -122,8 +122,43 @@ func TestToEscapingScheme(t *testing.T) {
 		},
 	}
 	for _, test := range tests {
-		if test.format.ToEscapingScheme() != test.expected {
-			t.Errorf("expected %v got %v", test.expected, test.format.ToEscapingScheme())
-		}
+		require.Equal(t, test.expected, test.format.ToEscapingScheme())
+	}
+}
+
+func TestWithEscapingScheme(t *testing.T) {
+	tests := []struct {
+		name     string
+		format   Format
+		scheme   model.EscapingScheme
+		expected string
+	}{
+		{
+			name:     "no existing term, append one",
+			format:   FmtProtoCompact,
+			scheme:   model.DotsEscaping,
+			expected: "application/vnd.google.protobuf; proto=io.prometheus.client.MetricFamily; encoding=compact-text; escaping=dots",
+		},
+		{
+			name:     "existing term at end, replace",
+			format:   "application/openmetrics-text; version=1.0.0; charset=utf-8; escaping=underscores",
+			scheme:   model.ValueEncodingEscaping,
+			expected: "application/openmetrics-text; version=1.0.0; charset=utf-8; escaping=values",
+		},
+		{
+			name:     "existing term in middle, replace",
+			format:   "application/openmetrics-text; escaping=dots; version=1.0.0; charset=utf-8; ",
+			scheme:   model.NoEscaping,
+			expected: "application/openmetrics-text; version=1.0.0; charset=utf-8; escaping=allow-utf-8",
+		},
+		{
+			name:     "multiple existing terms removed",
+			format:   "application/openmetrics-text; escaping=dots; version=1.0.0; charset=utf-8; escaping=allow-utf-8",
+			scheme:   model.ValueEncodingEscaping,
+			expected: "application/openmetrics-text; version=1.0.0; charset=utf-8; escaping=values",
+		},
+	}
+	for _, test := range tests {
+		require.Equal(t, test.expected, string(test.format.WithEscapingScheme(test.scheme)))
 	}
 }