From 726179cf0b3d4ce521dd170913503c05e82d1b46 Mon Sep 17 00:00:00 2001
From: Frank Jogeleit <fj@move-elevator.de>
Date: Thu, 10 Jun 2021 12:09:17 +0200
Subject: [PATCH] Add new APIs for Liveness and Readiness Probes

---
 pkg/api/handler.go      | 27 ++++++++++++++
 pkg/api/handler_test.go | 82 +++++++++++++++++++++++++++++++++++++++++
 pkg/api/server.go       |  2 +
 3 files changed, 111 insertions(+)

diff --git a/pkg/api/handler.go b/pkg/api/handler.go
index b82f43e..6da82a3 100644
--- a/pkg/api/handler.go
+++ b/pkg/api/handler.go
@@ -27,3 +27,30 @@ func PolicyHandler(s *kyverno.PolicyStore) http.HandlerFunc {
 		}
 	}
 }
+
+// HealthzHandler for the Liveness REST API
+func HealthzHandler(s *kyverno.PolicyStore) http.HandlerFunc {
+	return func(w http.ResponseWriter, req *http.Request) {
+		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
+
+		policies := s.List()
+		if len(policies) == 0 {
+			w.WriteHeader(http.StatusServiceUnavailable)
+			fmt.Fprint(w, `{ "error": "No Policies found" }`)
+
+			return
+		}
+
+		w.WriteHeader(http.StatusOK)
+		fmt.Fprint(w, "{}")
+	}
+}
+
+// ReadyHandler for the Readiness REST API
+func ReadyHandler() http.HandlerFunc {
+	return func(w http.ResponseWriter, req *http.Request) {
+		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
+		w.WriteHeader(http.StatusOK)
+		fmt.Fprint(w, "{}")
+	}
+}
diff --git a/pkg/api/handler_test.go b/pkg/api/handler_test.go
index dce3374..59d60c7 100644
--- a/pkg/api/handler_test.go
+++ b/pkg/api/handler_test.go
@@ -69,3 +69,85 @@ func Test_PolicyReportAPI(t *testing.T) {
 		}
 	})
 }
+
+func Test_HealthzAPI(t *testing.T) {
+	t.Run("No Policies Respose", func(t *testing.T) {
+		req, err := http.NewRequest("GET", "/healthz", nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		rr := httptest.NewRecorder()
+		handler := http.HandlerFunc(api.HealthzHandler(kyverno.NewPolicyStore()))
+
+		handler.ServeHTTP(rr, req)
+
+		if status := rr.Code; status != http.StatusServiceUnavailable {
+			t.Errorf("handler returned wrong status code: got %v want %v", status, http.StatusOK)
+		}
+
+		expected := `{ "error": "No Policies found" }`
+		if rr.Body.String() != expected {
+			t.Errorf("handler returned unexpected body: got %v want %v", rr.Body.String(), expected)
+		}
+	})
+	t.Run("Success Respose", func(t *testing.T) {
+		req, err := http.NewRequest("GET", "/healthz", nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		result := kyverno.Rule{
+			ValidateMessage: "validation error: requests and limits required. Rule autogen-check-for-requests-and-limits failed at path /spec/template/spec/containers/0/resources/requests/",
+			Name:            "autogen-check-for-requests-and-limits",
+		}
+
+		policy := kyverno.Policy{
+			Kind:              "Policy",
+			Name:              "require-ressources",
+			Namespace:         "test",
+			Rules:             []kyverno.Rule{result},
+			CreationTimestamp: time.Now(),
+		}
+
+		store := kyverno.NewPolicyStore()
+		store.Add(policy)
+
+		rr := httptest.NewRecorder()
+		handler := http.HandlerFunc(api.HealthzHandler(store))
+
+		handler.ServeHTTP(rr, req)
+
+		if status := rr.Code; status != http.StatusOK {
+			t.Errorf("handler returned wrong status code: got %v want %v", status, http.StatusOK)
+		}
+
+		expected := `{}`
+		if !strings.Contains(rr.Body.String(), expected) {
+			t.Errorf("handler returned unexpected body: got %v want %v", rr.Body.String(), expected)
+		}
+	})
+}
+
+func Test_ReadyAPI(t *testing.T) {
+	t.Run("Success Respose", func(t *testing.T) {
+		req, err := http.NewRequest("GET", "/healthz", nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		rr := httptest.NewRecorder()
+		handler := http.HandlerFunc(api.ReadyHandler())
+
+		handler.ServeHTTP(rr, req)
+
+		if status := rr.Code; status != http.StatusOK {
+			t.Errorf("handler returned wrong status code: got %v want %v", status, http.StatusOK)
+		}
+
+		expected := `{}`
+		if !strings.Contains(rr.Body.String(), expected) {
+			t.Errorf("handler returned unexpected body: got %v want %v", rr.Body.String(), expected)
+		}
+	})
+}
diff --git a/pkg/api/server.go b/pkg/api/server.go
index be14623..4f97662 100644
--- a/pkg/api/server.go
+++ b/pkg/api/server.go
@@ -21,6 +21,8 @@ type httpServer struct {
 
 func (s *httpServer) registerHandler() {
 	s.mux.HandleFunc("/policies", Gzip(PolicyHandler(s.store)))
+	s.mux.HandleFunc("/healthz", HealthzHandler(s.store))
+	s.mux.HandleFunc("/ready", ReadyHandler())
 }
 
 func (s *httpServer) Start() error {