package agent import ( "context" "crypto/tls" "crypto/x509" "encoding/json" "fmt" "net/http" "net/http/httptest" "os" "path/filepath" "testing" "time" "crypto/x509/pkix" "git.dws.rip/dubey/kat/internal/pki" ) 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) // Set the server certificate to use the test node name as CN // to match what our test agent will expect server.TLS.Certificates = []tls.Certificate{ { Certificate: [][]byte{[]byte("test-cert")}, PrivateKey: nil, Leaf: &x509.Certificate{ Subject: pkix.Name{ CommonName: "leader.kat.cluster.local", }, }, }, } // 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") }