Merge pull request #12419 from chlins/feat/p2p-preheat-healthcheck-controller

feat(preheat): add healthcheck methods for p2p preheat controller

src/controller/p2p/preheat/controller.go

@@ -2,9 +2,9 @@ package preheat
 
 import (
 	"context"
-	"errors"
 	"time"
 
+	"github.com/goharbor/harbor/src/lib/errors"
 	"github.com/goharbor/harbor/src/lib/q"
 	"github.com/goharbor/harbor/src/pkg/p2p/preheat/instance"
 	policyModels "github.com/goharbor/harbor/src/pkg/p2p/preheat/models/policy"
@@ -108,6 +108,8 @@ type Controller interface {
 	ListPolicies(ctx context.Context, query *q.Query) ([]*policyModels.Schema, error)
 	// ListPoliciesByProject lists policies by project.
 	ListPoliciesByProject(ctx context.Context, project int64, query *q.Query) ([]*policyModels.Schema, error)
+	// CheckHealth checks the instance health, for test connection
+	CheckHealth(ctx context.Context, instance *providerModels.Instance) error
 }
 
 var _ Controller = (*controller)(nil)
@@ -236,3 +238,33 @@ func (c *controller) ListPolicies(ctx context.Context, query *q.Query) ([]*polic
 func (c *controller) ListPoliciesByProject(ctx context.Context, project int64, query *q.Query) ([]*policyModels.Schema, error) {
 	return c.pManager.ListPoliciesByProject(ctx, project, query)
 }
+
+// CheckHealth checks the instance health, for test connection
+func (c *controller) CheckHealth(ctx context.Context, instance *providerModels.Instance) error {
+	if instance == nil {
+		return errors.New("instance can not be nil")
+	}
+
+	fac, ok := provider.GetProvider(instance.Vendor)
+	if !ok {
+		return errors.Errorf("no driver registered for provider %s", instance.Vendor)
+	}
+
+	// Construct driver
+	driver, err := fac(instance)
+	if err != nil {
+		return err
+	}
+
+	// Check health
+	h, err := driver.GetHealth()
+	if err != nil {
+		return err
+	}
+
+	if h.Status != provider.DriverStatusHealthy {
+		return errors.Errorf("preheat provider instance %s-%s:%s is not healthy", instance.Vendor, instance.Name, instance.Endpoint)
+	}
+
+	return nil
+}

src/controller/p2p/preheat/controllor_test.go

@@ -3,8 +3,12 @@ package preheat
 import (
 	"context"
 	"errors"
+	"net/http"
+	"net/http/httptest"
 	"testing"
 
+	"github.com/goharbor/harbor/src/pkg/p2p/preheat/provider/auth"
+
 	"github.com/goharbor/harbor/src/core/config"
 	"github.com/goharbor/harbor/src/pkg/p2p/preheat/models/policy"
 	providerModel "github.com/goharbor/harbor/src/pkg/p2p/preheat/models/provider"
@@ -18,10 +22,11 @@ import (
 
 type preheatSuite struct {
 	suite.Suite
-	ctx             context.Context
-	controller      Controller
-	fakeInstanceMgr *instance.FakeManager
-	fakePolicyMgr   *pmocks.FakeManager
+	ctx                context.Context
+	controller         Controller
+	fakeInstanceMgr    *instance.FakeManager
+	fakePolicyMgr      *pmocks.FakeManager
+	mockInstanceServer *httptest.Server
 }
 
 func TestPreheatSuite(t *testing.T) {
@@ -70,6 +75,25 @@ func (s *preheatSuite) SetupSuite() {
 		Endpoint: "http://localhost",
 	}, nil)
 	s.fakeInstanceMgr.On("Get", mock.Anything, int64(0)).Return(nil, errors.New("not found"))
+
+	// mock server for check health
+	s.mockInstanceServer = httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.RequestURI {
+		case "/_ping":
+			if r.Method != http.MethodGet {
+				w.WriteHeader(http.StatusNotImplemented)
+				return
+			}
+
+			w.WriteHeader(http.StatusOK)
+		}
+	}))
+	s.mockInstanceServer.Start()
+}
+
+// TearDownSuite clears the env.
+func (s *preheatSuite) TearDownSuite() {
+	s.mockInstanceServer.Close()
 }
 
 func (s *preheatSuite) TestGetAvailableProviders() {
@@ -213,3 +237,40 @@ func (s *preheatSuite) TestListPoliciesByProject() {
 	s.NoError(err)
 	s.NotNil(p)
 }
+
+func (s *preheatSuite) TestCheckHealth() {
+	// if instance is nil
+	var instance *providerModel.Instance
+	err := s.controller.CheckHealth(s.ctx, instance)
+	s.Error(err)
+
+	// unknown vendor
+	instance = &providerModel.Instance{
+		ID:       1,
+		Name:     "test-instance",
+		Vendor:   "unknown",
+		Endpoint: "http://127.0.0.1",
+		AuthMode: auth.AuthModeNone,
+		Enabled:  true,
+		Default:  true,
+		Insecure: true,
+		Status:   "Unknown",
+	}
+	err = s.controller.CheckHealth(s.ctx, instance)
+	s.Error(err)
+
+	// health
+	instance = &providerModel.Instance{
+		ID:       1,
+		Name:     "test-instance",
+		Vendor:   provider.DriverDragonfly,
+		Endpoint: s.mockInstanceServer.URL,
+		AuthMode: auth.AuthModeNone,
+		Enabled:  true,
+		Default:  true,
+		Insecure: true,
+		Status:   "Unknown",
+	}
+	err = s.controller.CheckHealth(s.ctx, instance)
+	s.NoError(err)
+}