138 lines
3.8 KiB
Go
138 lines
3.8 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"crypto/tls"
|
|
"crypto/x509"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.dws.rip/dubey/kat/internal/pki"
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
func TestAgentHeartbeat(t *testing.T) {
|
|
// Create temporary directory for test PKI files
|
|
tempDir, err := os.MkdirTemp("", "kat-test-agent-*")
|
|
if err != nil {
|
|
t.Fatalf("Failed to create temp directory: %v", err)
|
|
}
|
|
defer os.RemoveAll(tempDir)
|
|
|
|
// Generate CA for testing
|
|
pkiDir := filepath.Join(tempDir, "pki")
|
|
caKeyPath := filepath.Join(pkiDir, "ca.key")
|
|
caCertPath := filepath.Join(pkiDir, "ca.crt")
|
|
err = pki.GenerateCA(pkiDir, caKeyPath, caCertPath)
|
|
if err != nil {
|
|
t.Fatalf("Failed to generate test CA: %v", err)
|
|
}
|
|
|
|
// Generate node certificate
|
|
nodeKeyPath := filepath.Join(pkiDir, "node.key")
|
|
nodeCSRPath := filepath.Join(pkiDir, "node.csr")
|
|
nodeCertPath := filepath.Join(pkiDir, "node.crt")
|
|
err = pki.GenerateCertificateRequest("test-node", nodeKeyPath, nodeCSRPath)
|
|
if err != nil {
|
|
t.Fatalf("Failed to generate node key and CSR: %v", err)
|
|
}
|
|
err = pki.SignCertificateRequest(caKeyPath, caCertPath, nodeCSRPath, nodeCertPath, 24*time.Hour)
|
|
if err != nil {
|
|
t.Fatalf("Failed to sign node CSR: %v", err)
|
|
}
|
|
|
|
// Create a test server that requires client certificates
|
|
server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
// Verify the request path
|
|
if r.URL.Path != "/v1alpha1/nodes/test-node/status" {
|
|
t.Errorf("Expected path /v1alpha1/nodes/test-node/status, got %s", r.URL.Path)
|
|
http.Error(w, "Invalid path", http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
// Verify the request method
|
|
if r.Method != "POST" {
|
|
t.Errorf("Expected method POST, got %s", r.Method)
|
|
http.Error(w, "Invalid method", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
|
|
// Parse the request body
|
|
var status NodeStatus
|
|
decoder := json.NewDecoder(r.Body)
|
|
if err := decoder.Decode(&status); err != nil {
|
|
t.Errorf("Failed to decode request body: %v", err)
|
|
http.Error(w, "Invalid request body", http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
// Verify the node name
|
|
if status.NodeName != "test-node" {
|
|
t.Errorf("Expected node name test-node, got %s", status.NodeName)
|
|
http.Error(w, "Invalid node name", http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
// Verify that resources are present
|
|
if status.Resources.Capacity.CPU == "" || status.Resources.Capacity.Memory == "" {
|
|
t.Errorf("Missing resource capacity information")
|
|
http.Error(w, "Missing resource information", http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
// Return success
|
|
w.WriteHeader(http.StatusOK)
|
|
}))
|
|
defer server.Close()
|
|
|
|
// Configure the server to require client certificates
|
|
server.TLS.ClientAuth = tls.RequireAndVerifyClientCert
|
|
server.TLS.ClientCAs = x509.NewCertPool()
|
|
caCertData, err := os.ReadFile(caCertPath)
|
|
if err != nil {
|
|
t.Fatalf("Failed to read CA certificate: %v", err)
|
|
}
|
|
server.TLS.ClientCAs.AppendCertsFromPEM(caCertData)
|
|
|
|
// Extract the host:port from the server URL
|
|
serverURL := server.URL
|
|
hostPort := serverURL[8:] // Remove "https://" prefix
|
|
|
|
// Create an agent
|
|
agent, err := NewAgent("test-node", "test-uid", hostPort, "192.168.1.100", pkiDir, 1)
|
|
if err != nil {
|
|
t.Fatalf("Failed to create agent: %v", err)
|
|
}
|
|
|
|
// Setup mTLS client
|
|
err = agent.SetupMTLSClient()
|
|
if err != nil {
|
|
t.Fatalf("Failed to setup mTLS client: %v", err)
|
|
}
|
|
|
|
// Create a context with timeout
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
// Start heartbeat
|
|
err = agent.StartHeartbeat(ctx)
|
|
if err != nil {
|
|
t.Fatalf("Failed to start heartbeat: %v", err)
|
|
}
|
|
|
|
// Wait for at least one heartbeat
|
|
time.Sleep(2 * time.Second)
|
|
|
|
// Stop heartbeat
|
|
agent.StopHeartbeat()
|
|
|
|
// Test passed if we got here without errors
|
|
fmt.Println("Agent heartbeat test passed")
|
|
}
|