From 9f265beb2bbfb20db6130338ff8f481b0815f8fe Mon Sep 17 00:00:00 2001
From: okozachenko <okozachenko@vexxhost.com>
Date: Thu, 23 Jul 2020 17:21:50 +0300
Subject: [PATCH] Fix the logic to check compute node and add unit test for it

Change-Id: Icb7691c1f520a357f214b8cae19a748c278b22c9
---
 go.mod        |  2 ++
 go.sum        |  2 ++
 main.go       | 18 ++++++++++++------
 units_test.go | 33 +++++++++++++++++++++++++++++++++
 4 files changed, 49 insertions(+), 6 deletions(-)
 create mode 100644 units_test.go

diff --git a/go.mod b/go.mod
index e9677e7..b48ccbe 100644
--- a/go.mod
+++ b/go.mod
@@ -3,12 +3,14 @@ module github.com/vexxhost/node-labeler
 go 1.13
 
 require (
+	github.com/avast/retry-go v2.6.0+incompatible
 	github.com/evanphx/json-patch v4.5.0+incompatible
 	github.com/googleapis/gnostic v0.3.1 // indirect
 	github.com/gosimple/slug v1.7.0
 	github.com/jaypipes/ghw v0.5.0
 	github.com/prometheus/common v0.7.0
 	github.com/rainycape/unidecode v0.0.0-20150907023854-cb7f23ec59be // indirect
+	github.com/stretchr/testify v1.3.0
 	go.uber.org/atomic v1.4.0 // indirect
 	go.uber.org/multierr v1.1.0 // indirect
 	go.uber.org/zap v1.10.0
diff --git a/go.sum b/go.sum
index 65c138a..58a49c9 100644
--- a/go.sum
+++ b/go.sum
@@ -9,6 +9,8 @@ github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuy
 github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
 github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4 h1:Hs82Z41s6SdL1CELW+XaDYmOH4hkBN4/N9og/AsOv7E=
 github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
+github.com/avast/retry-go v2.6.0+incompatible h1:FelcMrm7Bxacr1/RM8+/eqkDkmVN7tjlsy51dOzB3LI=
+github.com/avast/retry-go v2.6.0+incompatible/go.mod h1:XtSnn+n/sHqQIpZ10K1qAevBhOOCWBLXXy3hyiqqBrY=
 github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24TaqPxmxbtue+5NUziq4I4S80YR8gNf3Q=
 github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+CedLV8=
 github.com/davecgh/go-spew v0.0.0-20151105211317-5215b55f46b2/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
diff --git a/main.go b/main.go
index 54150a6..b0c96bf 100644
--- a/main.go
+++ b/main.go
@@ -64,7 +64,7 @@ func main() {
 			"node.vexxhost.com/product": slug.Make(product.Name),
 		}
 
-		if !checkComputeNode() {
+		if checkComputeNode() {
 			labels["node-role.openstack.org"] = "compute"
 		}
 
@@ -94,17 +94,23 @@ func main() {
 // checkComputeNode checks if the current node is a compute node.
 func checkComputeNode() bool {
 
+	hostname, err := os.Hostname()
+	if err != nil {
+		log.Fatal(err.Error())
+	}
+	// Now there is only one check logic for now using hostname
+	return CheckComputeSubstring(hostname)
+}
+
+// checkComputeSubstring checks if the string includes compute substring.
+func CheckComputeSubstring(label string) bool {
 	compute_tag := []string{
 		"kvm",
 		"compute",
 	}
 
-	hostname, err := os.Hostname()
-	if err != nil {
-		log.Fatal(err.Error())
-	}
 	for _, sub := range compute_tag {
-		if strings.Contains(hostname, sub) {
+		if strings.Contains(label, sub) {
 			return true
 		}
 	}
diff --git a/units_test.go b/units_test.go
new file mode 100644
index 0000000..85dce64
--- /dev/null
+++ b/units_test.go
@@ -0,0 +1,33 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func assertCheckComputeSubstring(t *testing.T, label string, expected bool) {
+	res := CheckComputeSubstring(label)
+	assert.Equal(t, expected, res)
+}
+
+func TestCheckComputeSubstringWithkvmHostname(t *testing.T) {
+	hostname := string("kvm10.specterops.iad1.vexxhost.net")
+	expected := bool(true)
+	assertCheckComputeSubstring(t, hostname, expected)
+
+}
+
+func TestCheckComputeSubstringWithcomputeHostname(t *testing.T) {
+	hostname := string("compute.specterops.iad1.vexxhost.net")
+	expected := bool(true)
+	assertCheckComputeSubstring(t, hostname, expected)
+
+}
+
+func TestCheckComputeSubstringWithNoneComputeHostname(t *testing.T) {
+	hostname := string("ctl1.specterops.iad1.vexxhost.net")
+	expected := bool(false)
+	assertCheckComputeSubstring(t, hostname, expected)
+
+}