From e5d3ecb5a85265f52551dfe8fb3f0448fc5df97b Mon Sep 17 00:00:00 2001
From: Vicente Zepeda Mas <vzepedam@redhat.com>
Date: Thu, 4 Jul 2024 17:22:50 +0200
Subject: [PATCH] Add CISystem and BuildURL tracking

Signed-off-by: Vicente Zepeda Mas <vzepedam@redhat.com>
---
 pkg/elastic/document.go     | 10 +++++----
 pkg/elastic/elastic.go      | 25 +++++++++++++++++++++
 pkg/elastic/elastic_test.go | 43 +++++++++++++++++++++++++++++++++++++
 3 files changed, 74 insertions(+), 4 deletions(-)
 create mode 100644 pkg/elastic/elastic_test.go

diff --git a/pkg/elastic/document.go b/pkg/elastic/document.go
index 831a1ba..8ce411e 100644
--- a/pkg/elastic/document.go
+++ b/pkg/elastic/document.go
@@ -5,8 +5,10 @@ import (
 )
 
 type doc struct {
-	Uuid    string         `json:"uuid"`
-	Version string         `json:"version"`
-	Attack  string         `json:"attack"`
-	Metrics vegeta.Metrics `json:"metrics"`
+	Uuid     string         `json:"uuid"`
+	Version  string         `json:"version"`
+	Attack   string         `json:"attack"`
+	BuildURL string         `json:"buildUrl"`
+	CiSystem string         `json:"ciSystem"`
+	Metrics  vegeta.Metrics `json:"metrics"`
 }
diff --git a/pkg/elastic/elastic.go b/pkg/elastic/elastic.go
index 9b5d60e..901c180 100644
--- a/pkg/elastic/elastic.go
+++ b/pkg/elastic/elastic.go
@@ -3,6 +3,7 @@ package elastic
 import (
 	"context"
 	"fmt"
+	"os"
 	"strings"
 
 	"github.com/cloud-bulldozer/go-commons/indexers"
@@ -48,12 +49,16 @@ func IndexFile(ctx context.Context, testID string, version string, attack string
 		logger.Error(ctx, "obtaining indexer: %s", err)
 	}
 
+	env := extractENVContent()
+
 	var errors string
 	_doc := doc{}
 	_doc.Metrics = metrics
 	_doc.Uuid = testID
 	_doc.Version = version
 	_doc.Attack = attack
+	_doc.BuildURL = env["buildUrl"]
+	_doc.CiSystem = env["CiSystem"]
 
 	resp, err := (*indexer).Index([]interface{}{_doc}, indexers.IndexingOpts{
 		MetricName: strings.Join([]string{testID, attack}, "-"),
@@ -68,3 +73,23 @@ func IndexFile(ctx context.Context, testID string, version string, attack string
 	}
 	return nil
 }
+
+func extractENVContent() map[string]string {
+	result := make(map[string]string)
+	buildURL := os.Getenv("BUILD_URL")
+	if buildURL != "" {
+		result["buildUrl"] = buildURL
+	} else {
+		result["buildUrl"] = "Manual run"
+	}
+	dagID := os.Getenv("AIRFLOW_CTX_DAG_ID")
+	jenkinsURL := os.Getenv("JENKINS_URL")
+	if jenkinsURL != "" {
+		result["CiSystem"] = "Jenkins"
+	} else if dagID != "" {
+		result["CiSystem"] = "Airflow"
+	} else {
+		result["CiSystem"] = "Local"
+	}
+	return result
+}
diff --git a/pkg/elastic/elastic_test.go b/pkg/elastic/elastic_test.go
new file mode 100644
index 0000000..83afa96
--- /dev/null
+++ b/pkg/elastic/elastic_test.go
@@ -0,0 +1,43 @@
+package elastic
+
+import (
+	"os"
+	"reflect"
+	"testing"
+)
+
+func Test_extractENVContent(t *testing.T) {
+	tests := []struct {
+		name string
+		envs map[string]string
+		want map[string]string
+	}{
+		{"manual_run", map[string]string{}, map[string]string{"buildUrl": "Manual run", "CiSystem": "Local"}},
+		{"has_build_url", map[string]string{"BUILD_URL": "https://test.com"}, map[string]string{"buildUrl": "https://test.com", "CiSystem": "Local"}},
+		{"jenkins_run", map[string]string{"BUILD_URL": "https://test.com", "JENKINS_URL": "http://jenkins.dev"}, map[string]string{"buildUrl": "https://test.com", "CiSystem": "Jenkins"}},
+		{"airflow_run", map[string]string{"BUILD_URL": "https://test.com", "AIRFLOW_CTX_DAG_ID": "ocm-dag"}, map[string]string{"buildUrl": "https://test.com", "CiSystem": "Airflow"}},
+		{"jenkins_run", map[string]string{"BUILD_URL": "https://test.com", "JENKINS_URL": "http://jenkins.dev", "AIRFLOW_CTX_DAG_ID": "ocm-dag"}, map[string]string{"buildUrl": "https://test.com", "CiSystem": "Jenkins"}},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			for key, v := range tt.envs {
+				os.Setenv(key, v)
+			}
+			if got := extractENVContent(); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("extractENVContent() = %v, want %v", got, tt.want)
+			}
+			for key := range tt.envs {
+				os.Unsetenv(key)
+			}
+		})
+	}
+}
+
+// buildURL := os.Getenv("BUILD_URL")
+// if buildURL != "" {
+// 	result["buildUrl"] = buildURL
+// } else {
+// 	result["buildUrl"] = "Manual run"
+// }
+// dagID := os.Getenv("AIRFLOW_CTX_DAG_ID")
+// jenkinsURL := os.Getenv("JENKINS_URL")