From c2ed658aaae326ccbcdfa75078a9366736ae5e18 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adam=20=C5=A0trauch?= <cx@initd.cz>
Date: Mon, 21 Feb 2022 10:20:06 +0100
Subject: [PATCH] Multiple prometheus services support

* Prometheus export test
* Version bump to 1.5
---
 Makefile                  |  2 +-
 README.md                 |  2 ++
 daemon/prometheus.go      | 23 ++++++++++---------
 daemon/prometheus_test.go | 47 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 63 insertions(+), 11 deletions(-)
 create mode 100644 daemon/prometheus_test.go

diff --git a/Makefile b/Makefile
index 6f5f9a4..2dbd1c2 100644
--- a/Makefile
+++ b/Makefile
@@ -1,4 +1,4 @@
-VERSION=1.4
+VERSION=1.5
 
 
 .PHONY: all
diff --git a/README.md b/README.md
index bd9d38a..60f0ac4 100644
--- a/README.md
+++ b/README.md
@@ -137,6 +137,8 @@ Let's check this:
 
 If you set port to *-* lobby daemon omits port entirely from the output.
 
+There can be multiple `host` labels but only one `port` label and all prometheus labels (last line) will be common for all hosts labels. If port is omitted then default 9100 is used or port can also be part of the host label.
+
 When you open URL http://localhost:1313/v1/prometheus/nodeexporter it returns this:
 
 ```json
diff --git a/daemon/prometheus.go b/daemon/prometheus.go
index 78b6ea0..808f650 100644
--- a/daemon/prometheus.go
+++ b/daemon/prometheus.go
@@ -34,10 +34,10 @@ func preparePrometheusOutput(name string, discoveries []server.Discovery) Promet
 
 	for _, discovery := range discoveries {
 		port := strconv.Itoa(int(config.NodeExporterPort))
-		host := discovery.Hostname
+		hosts := []string{}
 		var add bool // add to the prometheus output when there is at least one prometheus related label
 
-		labels := map[string]string{}
+		labels := map[string]string{} // These are prometheus labels, not Lobby's labels
 
 		for _, label := range discovery.FindLabelsByPrefix("prometheus:" + name + ":") {
 			trimmed := strings.TrimPrefix(label.String(), "prometheus:"+name+":")
@@ -46,7 +46,7 @@ func preparePrometheusOutput(name string, discoveries []server.Discovery) Promet
 				if parts[0] == "port" {
 					port = parts[1]
 				} else if parts[0] == "host" {
-					host = parts[1]
+					hosts = append(hosts, parts[1])
 				} else {
 					labels[parts[0]] = parts[1]
 				}
@@ -65,20 +65,23 @@ func preparePrometheusOutput(name string, discoveries []server.Discovery) Promet
 		}
 
 		if add {
-			// Omit port part if "-" is set
-			target := host + ":" + port
-			if port == "-" {
-				target = host
-			}
+			targets := []string{}
+			for _, host := range hosts {
+				// Omit port part if "-" is set or port is part of the host
+				target := host + ":" + port
+				if strings.Contains(host, ":") || port == "-" {
+					target = host
+				}
+				targets = append(targets, target)
 
+			}
 			service := PrometheusService{
-				Targets: []string{target},
+				Targets: targets,
 				Labels:  labels,
 			}
 
 			services = append(services, service)
 		}
-
 	}
 
 	return services
diff --git a/daemon/prometheus_test.go b/daemon/prometheus_test.go
new file mode 100644
index 0000000..cb567ef
--- /dev/null
+++ b/daemon/prometheus_test.go
@@ -0,0 +1,47 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/by-cx/lobby/server"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestPreparePrometheusOutput(t *testing.T) {
+	discoveries := []server.Discovery{
+		{
+			Hostname: "test.server",
+			Labels: server.Labels{
+				// test1
+				"prometheus:test1:label1:l1",
+				"prometheus:test1:host:srv1:1234",
+				"prometheus:test1:host:srv1:1235",
+				"prometheus:test1:host:srv1",
+				"prometheus:test1:label2:l2",
+				// test2
+				"prometheus:test2:host:srv2",
+				"prometheus:test2:host:srv2:1235",
+				"prometheus:test2:host:srv2:1236",
+				"prometheus:test2:host:srv2:1237",
+				"prometheus:test2:label1:l3",
+			},
+		},
+	}
+	services := preparePrometheusOutput("test1", discoveries)
+	assert.Equal(t, 1, len(services))
+	assert.Equal(t, 3, len(services[0].Targets))
+	assert.Contains(t, services[0].Targets, "srv1:9100")
+	assert.Contains(t, services[0].Targets, "srv1:1234")
+	assert.Contains(t, services[0].Targets, "srv1:1235")
+	assert.Contains(t, services[0].Labels["label1"], "l1")
+	assert.Contains(t, services[0].Labels["label2"], "l2")
+
+	services = preparePrometheusOutput("test2", discoveries)
+	assert.Equal(t, 1, len(services))
+	assert.Equal(t, 4, len(services[0].Targets))
+	assert.Contains(t, services[0].Targets, "srv2:9100")
+	assert.Contains(t, services[0].Targets, "srv2:1235")
+	assert.Contains(t, services[0].Targets, "srv2:1236")
+	assert.Contains(t, services[0].Targets, "srv2:1237")
+	assert.Contains(t, services[0].Labels["label1"], "l3")
+}