diff --git a/src/controller/p2p/preheat/controller.go b/src/controller/p2p/preheat/controller.go index ce9a275ac..2398b84bb 100644 --- a/src/controller/p2p/preheat/controller.go +++ b/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 +} diff --git a/src/controller/p2p/preheat/controllor_test.go b/src/controller/p2p/preheat/controllor_test.go index 4744982c9..54f907ffc 100644 --- a/src/controller/p2p/preheat/controllor_test.go +++ b/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) +}