From 6344ee46026690ead89fbc97ed78a036283fb90c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Wilson=20J=C3=BAnior?= <wilsonpjunior@gmail.com>
Date: Mon, 7 Nov 2022 11:29:06 -0300
Subject: [PATCH] Allow to define classes also on SLO file

---
 main.go             | 10 ++++++++--
 slo/classes.go      |  8 +++++---
 slo/classes_test.go |  6 +++---
 slo/slo.go          |  3 ++-
 4 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/main.go b/main.go
index f6d9492..55802b7 100644
--- a/main.go
+++ b/main.go
@@ -66,6 +66,12 @@ func main() {
 		log.Fatal(err)
 	}
 
+	if len(spec.Classes) > 0 && len(classesDefinition.Classes) > 0 {
+		log.Fatal("you can not define classes in slo and classes files")
+	} else if len(classesDefinition.Classes) > 0 {
+		spec.Classes = classesDefinition.Classes
+	}
+
 	ruleGroups := &rulefmt.RuleGroups{
 		Groups: []rulefmt.RuleGroup{},
 	}
@@ -82,7 +88,7 @@ func main() {
 		manifests := []monitoringv1.PrometheusRule{}
 		for _, slo := range spec.SLOS {
 			// try to use any slo class found
-			sloClass, err := classesDefinition.FindClass(slo.Class)
+			sloClass, err := spec.Classes.FindClass(slo.Class)
 			if err != nil {
 				log.Fatalf("Could not compile SLO: %q, err: %q", slo.Name, err.Error())
 			}
@@ -126,7 +132,7 @@ func main() {
 
 	for _, slo := range spec.SLOS {
 		// try to use any slo class found
-		sloClass, err := classesDefinition.FindClass(slo.Class)
+		sloClass, err := spec.Classes.FindClass(slo.Class)
 		if err != nil {
 			log.Fatalf("Could not compile SLO: %q, err: %q", slo.Name, err.Error())
 		}
diff --git a/slo/classes.go b/slo/classes.go
index 1d27c61..e36c51d 100644
--- a/slo/classes.go
+++ b/slo/classes.go
@@ -11,15 +11,17 @@ type Class struct {
 }
 
 type ClassesDefinition struct {
-	Classes []Class `yaml:"classes"`
+	Classes Classes `yaml:"classes"`
 }
 
+type Classes []Class
+
 // FindClass finds for a given name, if not found return an error
-func (c *ClassesDefinition) FindClass(name string) (*Class, error) {
+func (c Classes) FindClass(name string) (*Class, error) {
 	if name == "" {
 		return nil, nil
 	}
-	for _, class := range c.Classes {
+	for _, class := range c {
 		if class.Name == name {
 			return &class, nil
 		}
diff --git a/slo/classes_test.go b/slo/classes_test.go
index 82b7094..a8675e3 100644
--- a/slo/classes_test.go
+++ b/slo/classes_test.go
@@ -14,15 +14,15 @@ func TestFindClass(t *testing.T) {
 		},
 	}
 
-	class, err := definition.FindClass("HIGH")
+	class, err := definition.Classes.FindClass("HIGH")
 	assert.NoError(t, err)
 	assert.Equal(t, &definition.Classes[0], class)
 
-	class, err = definition.FindClass("NOTFOUND")
+	class, err = definition.Classes.FindClass("NOTFOUND")
 	assert.EqualError(t, err, "SLO class \"NOTFOUND\" is not found")
 	assert.Nil(t, class)
 
-	class, err = definition.FindClass("")
+	class, err = definition.Classes.FindClass("")
 	assert.Nil(t, err)
 	assert.Nil(t, class)
 }
diff --git a/slo/slo.go b/slo/slo.go
index ea0b3c0..f02c8f0 100644
--- a/slo/slo.go
+++ b/slo/slo.go
@@ -32,7 +32,8 @@ var quantiles = []struct {
 }
 
 type SLOSpec struct {
-	SLOS []SLO
+	SLOS    []SLO   `yaml:"slos"`
+	Classes Classes `yaml:"classes"`
 }
 
 type ExprBlock struct {