Merge pull request #50 from Cloudzero/cp-22509

CP-22509: Search All Namespaces for KSM and Improve Search Logic
Cloudzero · Oct 16, 2024 · 33d5ebb · 33d5ebb
2 parents c6d5429 + f53e788
commit 33d5ebb
Show file tree

Hide file tree

Showing 4 changed files with 126 additions and 11 deletions.
diff --git a/pkg/cmd/config/command.go b/pkg/cmd/config/command.go
@@ -63,7 +63,7 @@ func NewCommand(ctx context.Context) *cli.Command {
 						return err
 					}
 
-					kubeStateMetricsURL, err := k8s.GetKubeStateMetricsURL(ctx, clientset, namespace)
+					kubeStateMetricsURL, err := k8s.GetKubeStateMetricsURL(ctx, clientset)
 					if err != nil {
 						return err
 					}

diff --git a/pkg/cmd/config/command_test.go b/pkg/cmd/config/command_test.go
@@ -14,11 +14,11 @@ import (
 	"github.com/cloudzero/cloudzero-agent-validator/pkg/k8s"
 )
 
-func TestGenerate(t *testing.T) {
+func TestGenerateByName(t *testing.T) {
 	// Define the namespace to be used in the test
 	namespace := "test-namespace"
 
-	// Create a fake clientset with some services
+	// Create a fake clientset with a service named "kube-state-metrics"
 	clientset := fake.NewSimpleClientset(
 		&corev1.Service{
 			ObjectMeta: metav1.ObjectMeta{
@@ -36,7 +36,81 @@ func TestGenerate(t *testing.T) {
 	ctx, _ := context.WithCancel(context.Background())
 
 	// Fetch the Kube State Metrics URL
-	kubeStateMetricsURL, err := k8s.GetKubeStateMetricsURL(ctx, clientset, namespace)
+	kubeStateMetricsURL, err := k8s.GetKubeStateMetricsURL(ctx, clientset)
+	assert.NoError(t, err)
+
+	// Define the scrape config data
+	scrapeConfigData := config.ScrapeConfigData{
+		Targets:        []string{kubeStateMetricsURL},
+		ClusterName:    "test-cluster",
+		CloudAccountID: "123456789",
+		Region:         "us-west-2",
+		Host:           "test-host",
+		SecretPath:     "/etc/config/prometheus/secrets/",
+	}
+
+	// Generate the configuration content
+	configContent, err := config.Generate(scrapeConfigData)
+	assert.NoError(t, err)
+	assert.NotEmpty(t, configContent)
+
+	// Validate the dynamically populated values
+	assert.Contains(t, configContent, kubeStateMetricsURL)
+	assert.Contains(t, configContent, "cluster_name=test-cluster")
+	assert.Contains(t, configContent, "cloud_account_id=123456789")
+	assert.Contains(t, configContent, "region=us-west-2")
+	assert.Contains(t, configContent, "test-host")
+	assert.Contains(t, configContent, "/etc/config/prometheus/secrets/")
+
+	// Define the ConfigMap data
+	configMapData := map[string]string{
+		"prometheus.yml": configContent,
+	}
+
+	// Update the ConfigMap
+	err = k8s.UpdateConfigMap(ctx, clientset, namespace, "test-configmap", configMapData)
+	assert.NoError(t, err)
+
+	// Verify the ConfigMap was updated
+	updatedConfigMap, err := clientset.CoreV1().ConfigMaps(namespace).Get(ctx, "test-configmap", metav1.GetOptions{})
+	assert.NoError(t, err)
+	assert.Equal(t, configContent, updatedConfigMap.Data["prometheus.yml"])
+
+	// Clean up the output file if it exists
+	outputFile := "test_output.yml"
+	if _, err := os.Stat(outputFile); err == nil {
+		err = os.Remove(outputFile)
+		assert.NoError(t, err)
+	}
+}
+
+func TestGenerateByLabel(t *testing.T) {
+	// Define the namespace to be used in the test
+	namespace := "test-namespace"
+
+	// Create a fake clientset with a service having Helm-specific labels
+	clientset := fake.NewSimpleClientset(
+		&corev1.Service{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      "custom-service-name",
+				Namespace: namespace,
+				Labels: map[string]string{
+					"app.kubernetes.io/name": "kube-state-metrics",
+					"helm.sh/chart":          "kube-state-metrics-2.11.1",
+				},
+			},
+			Spec: corev1.ServiceSpec{
+				Ports: []corev1.ServicePort{
+					{Port: 8080},
+				},
+			},
+		},
+	)
+
+	ctx, _ := context.WithCancel(context.Background())
+
+	// Fetch the Kube State Metrics URL
+	kubeStateMetricsURL, err := k8s.GetKubeStateMetricsURL(ctx, clientset)
 	assert.NoError(t, err)
 
 	// Define the scrape config data

diff --git a/pkg/k8s/services.go b/pkg/k8s/services.go
@@ -53,9 +53,10 @@ func BuildKubeClient(kubeconfigPath string) (kubernetes.Interface, error) {
 	return clientset, nil
 }
 
-// GetKubeStateMetricsURL fetches the URL for the Kube State Metrics service
-func GetKubeStateMetricsURL(ctx context.Context, clientset kubernetes.Interface, namespace string) (string, error) {
-	services, err := clientset.CoreV1().Services(namespace).List(ctx, metav1.ListOptions{})
+// GetKubeStateMetricsURL fetches the URL for the Kube State Metrics service across all namespaces
+func GetKubeStateMetricsURL(ctx context.Context, clientset kubernetes.Interface) (string, error) {
+	// First, try to find the service by name
+	services, err := clientset.CoreV1().Services("").List(ctx, metav1.ListOptions{})
 	if err != nil {
 		return "", errors.Wrap(err, "listing services")
 	}
@@ -65,7 +66,17 @@ func GetKubeStateMetricsURL(ctx context.Context, clientset kubernetes.Interface,
 	for _, service := range services.Items {
 		if strings.Contains(service.Name, "kube-state-metrics") {
 			kubeStateMetricsURL = fmt.Sprintf("%s.%s.svc.cluster.local:%d", service.Name, service.Namespace, service.Spec.Ports[0].Port)
-			break
+			return kubeStateMetricsURL, nil
+		}
+	}
+
+	// If not found by name, check by labels
+	for _, service := range services.Items {
+		// Check for Helm-specific labels
+		if service.Labels["app.kubernetes.io/name"] == "kube-state-metrics" &&
+			service.Labels["helm.sh/chart"] != "" { // Ensure the service is managed by Helm
+			kubeStateMetricsURL = fmt.Sprintf("%s.%s.svc.cluster.local:%d", service.Name, service.Namespace, service.Spec.Ports[0].Port)
+			return kubeStateMetricsURL, nil
 		}
 	}
 

diff --git a/pkg/k8s/services_test.go b/pkg/k8s/services_test.go
@@ -13,12 +13,12 @@ import (
 	"github.com/cloudzero/cloudzero-agent-validator/pkg/k8s"
 )
 
-func TestGetKubeStateMetricsURL(t *testing.T) {
+func TestGetKubeStateMetricsURLByName(t *testing.T) {
 	clientset := fake.NewSimpleClientset()
 	ctx := context.TODO()
 	namespace := "test-namespace"
 
-	// Create a fake service in the test namespace
+	// Create a fake service in the test namespace with the name "kube-state-metrics"
 	_, err := clientset.CoreV1().Services(namespace).Create(ctx, &corev1.Service{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      "kube-state-metrics",
@@ -34,11 +34,41 @@ func TestGetKubeStateMetricsURL(t *testing.T) {
 	}, metav1.CreateOptions{})
 	assert.NoError(t, err)
 
-	kubeStateMetricsURL, err := k8s.GetKubeStateMetricsURL(ctx, clientset, namespace)
+	kubeStateMetricsURL, err := k8s.GetKubeStateMetricsURL(ctx, clientset)
 	assert.NoError(t, err)
 	assert.Contains(t, kubeStateMetricsURL, "kube-state-metrics")
 }
 
+func TestGetKubeStateMetricsURLByLabel(t *testing.T) {
+	clientset := fake.NewSimpleClientset()
+	ctx := context.TODO()
+	namespace := "test-namespace"
+
+	// Create a fake service in the test namespace with Helm-specific labels
+	_, err := clientset.CoreV1().Services(namespace).Create(ctx, &corev1.Service{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "custom-service-name",
+			Namespace: namespace,
+			Labels: map[string]string{
+				"app.kubernetes.io/name": "kube-state-metrics",
+				"helm.sh/chart":          "kube-state-metrics-2.11.1",
+			},
+		},
+		Spec: corev1.ServiceSpec{
+			Ports: []corev1.ServicePort{
+				{
+					Port: 8080,
+				},
+			},
+		},
+	}, metav1.CreateOptions{})
+	assert.NoError(t, err)
+
+	kubeStateMetricsURL, err := k8s.GetKubeStateMetricsURL(ctx, clientset)
+	assert.NoError(t, err)
+	assert.Contains(t, kubeStateMetricsURL, "custom-service-name")
+}
+
 func TestUpdateConfigMap(t *testing.T) {
 	clientset := fake.NewSimpleClientset()
 	ctx := context.TODO()