From d1f48ce043827c79b14bf36bd4fb5ad48fe83edc Mon Sep 17 00:00:00 2001
From: Steven Kang <skan070@gmail.com>
Date: Mon, 10 Feb 2025 10:45:43 +1300
Subject: [PATCH] feat: improve diagnostics stability - release 2.27 (#398)

---
 pkg/endpoints/utils.go          | 14 ++++++++
 pkg/endpoints/utils_test.go     | 58 +++++++++++++++++++++++++++++++++
 pkg/snapshot/docker_test.go     |  1 -
 pkg/snapshot/kubernetes_test.go | 43 ++++++++++++++++++++++++
 4 files changed, 115 insertions(+), 1 deletion(-)
 delete mode 100644 pkg/snapshot/docker_test.go

diff --git a/pkg/endpoints/utils.go b/pkg/endpoints/utils.go
index e32b56fea..4c704d86f 100644
--- a/pkg/endpoints/utils.go
+++ b/pkg/endpoints/utils.go
@@ -1,6 +1,7 @@
 package endpoints
 
 import (
+	"github.com/hashicorp/go-version"
 	portainer "github.com/portainer/portainer/api"
 )
 
@@ -15,6 +16,11 @@ func IsEdgeEndpoint(endpoint *portainer.Endpoint) bool {
 	return endpoint.Type == portainer.EdgeAgentOnDockerEnvironment || endpoint.Type == portainer.EdgeAgentOnKubernetesEnvironment
 }
 
+// IsStandardEdgeEndpoint returns true if this is a standard Edge endpoint and not in async mode on either Docker or Kubernetes
+func IsStandardEdgeEndpoint(endpoint *portainer.Endpoint) bool {
+	return (endpoint.Type == portainer.EdgeAgentOnDockerEnvironment || endpoint.Type == portainer.EdgeAgentOnKubernetesEnvironment) && !endpoint.Edge.AsyncMode
+}
+
 // IsAssociatedEdgeEndpoint returns true if the environment is an Edge environment
 // and has a set EdgeID and UserTrusted is true.
 func IsAssociatedEdgeEndpoint(endpoint *portainer.Endpoint) bool {
@@ -26,3 +32,11 @@ func IsAssociatedEdgeEndpoint(endpoint *portainer.Endpoint) bool {
 func HasDirectConnectivity(endpoint *portainer.Endpoint) bool {
 	return !IsEdgeEndpoint(endpoint) || (IsAssociatedEdgeEndpoint(endpoint) && !endpoint.Edge.AsyncMode)
 }
+
+// IsNewerThan225 returns true if the agent version is newer than 2.25.0
+// this is used to check if the agent is compatible with the new diagnostics feature
+func IsNewerThan225(agentVersion string) bool {
+	v1, _ := version.NewVersion(agentVersion)
+	v2, _ := version.NewVersion("2.25.0")
+	return v1.GreaterThanOrEqual(v2)
+}
diff --git a/pkg/endpoints/utils_test.go b/pkg/endpoints/utils_test.go
index 47ef2db8d..4231cf60a 100644
--- a/pkg/endpoints/utils_test.go
+++ b/pkg/endpoints/utils_test.go
@@ -202,3 +202,61 @@ func TestHasDirectConnectivity(t *testing.T) {
 		})
 	}
 }
+
+func TestIsStandardEdgeEndpoint(t *testing.T) {
+	tests := []struct {
+		name     string
+		endpoint *portainer.Endpoint
+		expected bool
+	}{
+		{
+			name: "StandardEdgeEndpoint",
+			endpoint: &portainer.Endpoint{
+				Type: portainer.EdgeAgentOnDockerEnvironment,
+				Edge: portainer.EnvironmentEdgeSettings{AsyncMode: false},
+			},
+			expected: true,
+		},
+		{
+			name: "AsyncEdgeEndpoint",
+			endpoint: &portainer.Endpoint{
+				Type: portainer.EdgeAgentOnDockerEnvironment,
+				Edge: portainer.EnvironmentEdgeSettings{AsyncMode: true},
+			},
+			expected: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := IsStandardEdgeEndpoint(tt.endpoint)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+func TestIsNewerThan225(t *testing.T) {
+	tests := []struct {
+		name     string
+		version  string
+		expected bool
+	}{
+		{
+			name:     "NewerThan225",
+			version:  "2.25.1",
+			expected: true,
+		},
+		{
+			name:     "OlderThan225",
+			version:  "2.24.0",
+			expected: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := IsNewerThan225(tt.version)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
diff --git a/pkg/snapshot/docker_test.go b/pkg/snapshot/docker_test.go
deleted file mode 100644
index 8df14bc39..000000000
--- a/pkg/snapshot/docker_test.go
+++ /dev/null
@@ -1 +0,0 @@
-package snapshot
diff --git a/pkg/snapshot/kubernetes_test.go b/pkg/snapshot/kubernetes_test.go
index 8df14bc39..142241a5c 100644
--- a/pkg/snapshot/kubernetes_test.go
+++ b/pkg/snapshot/kubernetes_test.go
@@ -1 +1,44 @@
 package snapshot
+
+import (
+	"context"
+	"testing"
+
+	portainer "github.com/portainer/portainer/api"
+	"github.com/stretchr/testify/require"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	kfake "k8s.io/client-go/kubernetes/fake"
+)
+
+func TestCreateKubernetesSnapshot(t *testing.T) {
+	cli := kfake.NewSimpleClientset()
+	kubernetesSnapshot := &portainer.KubernetesSnapshot{}
+
+	serverInfo, err := cli.Discovery().ServerVersion()
+	if err != nil {
+		t.Fatalf("error getting the kubernetesserver version: %v", err)
+	}
+
+	kubernetesSnapshot.KubernetesVersion = serverInfo.GitVersion
+	require.Equal(t, kubernetesSnapshot.KubernetesVersion, serverInfo.GitVersion)
+
+	nodeList, err := cli.CoreV1().Nodes().List(context.TODO(), metav1.ListOptions{})
+	if err != nil {
+		t.Fatalf("error listing kubernetes nodes: %v", err)
+	}
+
+	var totalCPUs, totalMemory int64
+	for _, node := range nodeList.Items {
+		totalCPUs += node.Status.Capacity.Cpu().Value()
+		totalMemory += node.Status.Capacity.Memory().Value()
+	}
+
+	kubernetesSnapshot.TotalCPU = totalCPUs
+	kubernetesSnapshot.TotalMemory = totalMemory
+	kubernetesSnapshot.NodeCount = len(nodeList.Items)
+	require.Equal(t, kubernetesSnapshot.TotalCPU, totalCPUs)
+	require.Equal(t, kubernetesSnapshot.TotalMemory, totalMemory)
+	require.Equal(t, kubernetesSnapshot.NodeCount, len(nodeList.Items))
+
+	t.Logf("Kubernetes snapshot: %+v", kubernetesSnapshot)
+}