From bc990582f525d733adf4c8b1dc6765713665836b Mon Sep 17 00:00:00 2001
From: Lukasz Mierzwa <lukasz@cloudflare.com>
Date: Tue, 27 Aug 2024 08:03:49 +0100
Subject: [PATCH] Fix benchmarks

---
 internal/config/config.go    | 9 ++++++++-
 internal/config/match.go     | 2 +-
 internal/config/rule_test.go | 9 +++++++++
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/internal/config/config.go b/internal/config/config.go
index abe59a2a..83cff768 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -83,7 +83,7 @@ func (cfg Config) String() string {
 func (cfg *Config) GetChecksForEntry(ctx context.Context, gen *PrometheusGenerator, entry discovery.Entry) []checks.RuleChecker {
 	enabled := []checks.RuleChecker{}
 
-	defaultMatch := []Match{{State: defaultMatchStates(ctx.Value(CommandKey).(ContextCommandVal))}}
+	defaultMatch := []Match{{State: defaultMatchStates(commandFromContext(ctx))}}
 	proms := gen.ServersForPath(entry.Path.Name)
 
 	if entry.PathError != nil || entry.Rule.Error.Err != nil {
@@ -229,3 +229,10 @@ func parseDuration(d string) (time.Duration, error) {
 	}
 	return time.Duration(mdur), nil
 }
+
+func commandFromContext(ctx context.Context) (cmd ContextCommandVal) {
+	if val := ctx.Value(CommandKey); val != nil {
+		cmd = val.(ContextCommandVal)
+	}
+	return cmd
+}
diff --git a/internal/config/match.go b/internal/config/match.go
index c7ac8b12..4a20a48b 100644
--- a/internal/config/match.go
+++ b/internal/config/match.go
@@ -105,7 +105,7 @@ func (m Match) validate(allowEmpty bool) error {
 }
 
 func (m Match) IsMatch(ctx context.Context, path string, e discovery.Entry) bool {
-	cmd := ctx.Value(CommandKey).(ContextCommandVal)
+	cmd := commandFromContext(ctx)
 
 	if m.Command != nil {
 		if cmd != *m.Command {
diff --git a/internal/config/rule_test.go b/internal/config/rule_test.go
index 7870bf5f..8ac270ac 100644
--- a/internal/config/rule_test.go
+++ b/internal/config/rule_test.go
@@ -665,6 +665,15 @@ func TestMatch(t *testing.T) {
 			},
 			isMatch: true,
 		},
+		{
+			path: "foo.yaml",
+			entry: discovery.Entry{
+				Rule:  parser.Rule{},
+				State: discovery.Noop,
+			},
+			match:   config.Match{},
+			isMatch: true,
+		},
 	}
 
 	for i, tc := range testCases {