feat: add p2p preheat controller policy releated logic

Signed-off-by: chlins <chlins.zhang@gmail.com>

src/controller/p2p/preheat/controller.go

@@ -6,9 +6,10 @@ import (
 	"time"
 
 	"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"
 	providerModels "github.com/goharbor/harbor/src/pkg/p2p/preheat/models/provider"
+	"github.com/goharbor/harbor/src/pkg/p2p/preheat/policy"
 	"github.com/goharbor/harbor/src/pkg/p2p/preheat/provider"
 )
 
@@ -85,6 +86,23 @@ type Controller interface {
 	// Any problems met, a non nil error will be returned
 	//
 	UpdateInstance(ctx context.Context, instance *providerModels.Instance, properties ...string) error
+
+	// do not provide another policy controller, mixed in preheat controller
+
+	// CountPolicy returns the total count of the policy.
+	CountPolicy(ctx context.Context, query *q.Query) (int64, error)
+	// CreatePolicy creates the policy.
+	CreatePolicy(ctx context.Context, schema *policyModels.Schema) (int64, error)
+	// GetPolicy gets the policy by id.
+	GetPolicy(ctx context.Context, id int64) (*policyModels.Schema, error)
+	// UpdatePolicy updates the policy.
+	UpdatePolicy(ctx context.Context, schema *policyModels.Schema, props ...string) error
+	// DeletePolicy deletes the policy by id.
+	DeletePolicy(ctx context.Context, id int64) error
+	// ListPolicies lists policies by query.
+	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)
 }
 
 var _ Controller = (*controller)(nil)
@@ -94,32 +112,35 @@ var _ Controller = (*controller)(nil)
 type controller struct {
 	// For instance
 	iManager instance.Manager
+	// For policy
+	pManager policy.Manager
 }
 
 // NewController is constructor of controller
 func NewController() Controller {
 	return &controller{
 		iManager: instance.Mgr,
+		pManager: policy.Mgr,
 	}
 }
 
 // GetAvailableProviders implements @Controller.GetAvailableProviders
-func (cc *controller) GetAvailableProviders() ([]*provider.Metadata, error) {
+func (c *controller) GetAvailableProviders() ([]*provider.Metadata, error) {
 	return provider.ListProviders()
 }
 
 // CountInstance implements @Controller.CountInstance
-func (cc *controller) CountInstance(ctx context.Context, query *q.Query) (int64, error) {
+func (c *controller) CountInstance(ctx context.Context, query *q.Query) (int64, error) {
-	return cc.iManager.Count(ctx, query)
+	return c.iManager.Count(ctx, query)
 }
 
-// List implements @Controller.ListInstance
+// ListInstance implements @Controller.ListInstance
-func (cc *controller) ListInstance(ctx context.Context, query *q.Query) ([]*providerModels.Instance, error) {
+func (c *controller) ListInstance(ctx context.Context, query *q.Query) ([]*providerModels.Instance, error) {
-	return cc.iManager.List(ctx, query)
+	return c.iManager.List(ctx, query)
 }
 
 // CreateInstance implements @Controller.CreateInstance
-func (cc *controller) CreateInstance(ctx context.Context, instance *providerModels.Instance) (int64, error) {
+func (c *controller) CreateInstance(ctx context.Context, instance *providerModels.Instance) (int64, error) {
 	if instance == nil {
 		return 0, errors.New("nil instance object provided")
 	}
@@ -130,7 +151,7 @@ func (cc *controller) CreateInstance(ctx context.Context, instance *providerMode
 			"endpoint": instance.Endpoint,
 		},
 	}
-	num, err := cc.iManager.Count(ctx, query)
+	num, err := c.iManager.Count(ctx, query)
 	if err != nil {
 		return 0, err
 	}
@@ -145,24 +166,59 @@ func (cc *controller) CreateInstance(ctx context.Context, instance *providerMode
 
 	instance.SetupTimestamp = time.Now().Unix()
 
-	return cc.iManager.Save(ctx, instance)
+	return c.iManager.Save(ctx, instance)
 }
 
-// Delete implements @Controller.Delete
+// DeleteInstance implements @Controller.Delete
-func (cc *controller) DeleteInstance(ctx context.Context, id int64) error {
+func (c *controller) DeleteInstance(ctx context.Context, id int64) error {
-	return cc.iManager.Delete(ctx, id)
+	return c.iManager.Delete(ctx, id)
 }
 
-// Update implements @Controller.Update
+// UpdateInstance implements @Controller.Update
-func (cc *controller) UpdateInstance(ctx context.Context, instance *providerModels.Instance, properties ...string) error {
+func (c *controller) UpdateInstance(ctx context.Context, instance *providerModels.Instance, properties ...string) error {
 	if len(properties) == 0 {
 		return errors.New("no properties provided to update")
 	}
 
-	return cc.iManager.Update(ctx, instance, properties...)
+	return c.iManager.Update(ctx, instance, properties...)
 }
 
-// Get implements @Controller.Get
+// GetInstance implements @Controller.Get
-func (cc *controller) GetInstance(ctx context.Context, id int64) (*providerModels.Instance, error) {
+func (c *controller) GetInstance(ctx context.Context, id int64) (*providerModels.Instance, error) {
-	return cc.iManager.Get(ctx, id)
+	return c.iManager.Get(ctx, id)
+}
+
+// CountPolicy returns the total count of the policy.
+func (c *controller) CountPolicy(ctx context.Context, query *q.Query) (int64, error) {
+	return c.pManager.Count(ctx, query)
+}
+
+// CreatePolicy creates the policy.
+func (c *controller) CreatePolicy(ctx context.Context, schema *policyModels.Schema) (int64, error) {
+	return c.pManager.Create(ctx, schema)
+}
+
+// GetPolicy gets the policy by id.
+func (c *controller) GetPolicy(ctx context.Context, id int64) (*policyModels.Schema, error) {
+	return c.pManager.Get(ctx, id)
+}
+
+// UpdatePolicy updates the policy.
+func (c *controller) UpdatePolicy(ctx context.Context, schema *policyModels.Schema, props ...string) error {
+	return c.pManager.Update(ctx, schema, props...)
+}
+
+// DeletePolicy deletes the policy by id.
+func (c *controller) DeletePolicy(ctx context.Context, id int64) error {
+	return c.pManager.Delete(ctx, id)
+}
+
+// ListPolicies lists policies by query.
+func (c *controller) ListPolicies(ctx context.Context, query *q.Query) ([]*policyModels.Schema, error) {
+	return c.pManager.ListPolicies(ctx, query)
+}
+
+// ListPoliciesByProject lists policies by project.
+func (c *controller) ListPoliciesByProject(ctx context.Context, project int64, query *q.Query) ([]*policyModels.Schema, error) {
+	return c.pManager.ListPoliciesByProject(ctx, project, query)
 }

src/controller/p2p/preheat/controllor_test.go

@@ -6,9 +6,11 @@ import (
 	"testing"
 
 	"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"
 	"github.com/goharbor/harbor/src/pkg/p2p/preheat/provider"
 	"github.com/goharbor/harbor/src/testing/pkg/p2p/preheat/instance"
+	pmocks "github.com/goharbor/harbor/src/testing/pkg/p2p/preheat/policy"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/suite"
@@ -16,24 +18,28 @@ import (
 
 type preheatSuite struct {
 	suite.Suite
-	ctx         context.Context
+	ctx             context.Context
-	controller  Controller
+	controller      Controller
-	fackManager *instance.FakeManager
+	fakeInstanceMgr *instance.FakeManager
+	fakePolicyMgr   *pmocks.FakeManager
 }
 
 func TestPreheatSuite(t *testing.T) {
 	t.Log("Start TestPreheatSuite")
-	fackManager := &instance.FakeManager{}
+	fakeInstanceMgr := &instance.FakeManager{}
+	fakePolicyMgr := &pmocks.FakeManager{}
 
 	var c = &controller{
-		iManager: fackManager,
+		iManager: fakeInstanceMgr,
+		pManager: fakePolicyMgr,
 	}
 	assert.NotNil(t, c)
 
 	suite.Run(t, &preheatSuite{
-		ctx:         context.Background(),
+		ctx:             context.Background(),
-		controller:  c,
+		controller:      c,
-		fackManager: fackManager,
+		fakeInstanceMgr: fakeInstanceMgr,
+		fakePolicyMgr:   fakePolicyMgr,
 	})
 }
 
@@ -45,7 +51,7 @@ func TestNewController(t *testing.T) {
 func (s *preheatSuite) SetupSuite() {
 	config.Init()
 
-	s.fackManager.On("List", mock.Anything, mock.Anything).Return([]*providerModel.Instance{
+	s.fakeInstanceMgr.On("List", mock.Anything, mock.Anything).Return([]*providerModel.Instance{
 		{
 			ID:       1,
 			Vendor:   "dragonfly",
@@ -54,16 +60,16 @@ func (s *preheatSuite) SetupSuite() {
 			Enabled:  true,
 		},
 	}, nil)
-	s.fackManager.On("Save", mock.Anything, mock.Anything).Return(int64(1), nil)
+	s.fakeInstanceMgr.On("Save", mock.Anything, mock.Anything).Return(int64(1), nil)
-	s.fackManager.On("Count", mock.Anything, &providerModel.Instance{Endpoint: "http://localhost"}).Return(int64(1), nil)
+	s.fakeInstanceMgr.On("Count", mock.Anything, &providerModel.Instance{Endpoint: "http://localhost"}).Return(int64(1), nil)
-	s.fackManager.On("Count", mock.Anything, mock.Anything).Return(int64(0), nil)
+	s.fakeInstanceMgr.On("Count", mock.Anything, mock.Anything).Return(int64(0), nil)
-	s.fackManager.On("Delete", mock.Anything, int64(1)).Return(nil)
+	s.fakeInstanceMgr.On("Delete", mock.Anything, int64(1)).Return(nil)
-	s.fackManager.On("Delete", mock.Anything, int64(0)).Return(errors.New("not found"))
+	s.fakeInstanceMgr.On("Delete", mock.Anything, int64(0)).Return(errors.New("not found"))
-	s.fackManager.On("Get", mock.Anything, int64(1)).Return(&providerModel.Instance{
+	s.fakeInstanceMgr.On("Get", mock.Anything, int64(1)).Return(&providerModel.Instance{
 		ID:       1,
 		Endpoint: "http://localhost",
 	}, nil)
-	s.fackManager.On("Get", mock.Anything, int64(0)).Return(nil, errors.New("not found"))
+	s.fakeInstanceMgr.On("Get", mock.Anything, int64(0)).Return(nil, errors.New("not found"))
 }
 
 func (s *preheatSuite) TestGetAvailableProviders() {
@@ -136,14 +142,14 @@ func (s *preheatSuite) TestDeleteInstance() {
 
 func (s *preheatSuite) TestUpdateInstance() {
 	// TODO: test update more
-	s.fackManager.On("Update", s.ctx, nil).Return(errors.New("no properties provided to update"))
+	s.fakeInstanceMgr.On("Update", s.ctx, nil).Return(errors.New("no properties provided to update"))
 	err := s.controller.UpdateInstance(s.ctx, nil)
 	s.Error(err)
 
 	err = s.controller.UpdateInstance(s.ctx, &providerModel.Instance{ID: 0})
 	s.Error(err)
 
-	s.fackManager.On("Update", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	s.fakeInstanceMgr.On("Update", mock.Anything, mock.Anything, mock.Anything).Return(nil)
 	err = s.controller.UpdateInstance(s.ctx, &providerModel.Instance{ID: 1}, "enabled")
 	s.NoError(err)
 }
@@ -153,3 +159,50 @@ func (s *preheatSuite) TestGetInstance() {
 	s.NoError(err)
 	s.NotNil(inst)
 }
+
+func (s *preheatSuite) TestCountPolicy() {
+	s.fakePolicyMgr.On("Count", s.ctx, mock.Anything).Return(int64(1), nil)
+	id, err := s.controller.CountPolicy(s.ctx, nil)
+	s.NoError(err)
+	s.Equal(int64(1), id)
+}
+
+func (s *preheatSuite) TestCreatePolicy() {
+	s.fakePolicyMgr.On("Create", s.ctx, mock.Anything).Return(int64(1), nil)
+	id, err := s.controller.CreatePolicy(s.ctx, nil)
+	s.NoError(err)
+	s.Equal(int64(1), id)
+}
+
+func (s *preheatSuite) TestGetPolicy() {
+	s.fakePolicyMgr.On("Get", s.ctx, int64(1)).Return(&policy.Schema{Name: "test"}, nil)
+	p, err := s.controller.GetPolicy(s.ctx, 1)
+	s.NoError(err)
+	s.Equal("test", p.Name)
+}
+
+func (s *preheatSuite) TestUpdatePolicy() {
+	s.fakePolicyMgr.On("Update", s.ctx, mock.Anything, mock.Anything).Return(nil)
+	err := s.controller.UpdateInstance(s.ctx, nil, "")
+	s.NoError(err)
+}
+
+func (s *preheatSuite) TestDeletePolicy() {
+	s.fakePolicyMgr.On("Delete", s.ctx, int64(1)).Return(nil)
+	err := s.controller.DeletePolicy(s.ctx, 1)
+	s.NoError(err)
+}
+
+func (s *preheatSuite) TestListPolicies() {
+	s.fakePolicyMgr.On("ListPolicies", s.ctx, mock.Anything).Return([]*policy.Schema{}, nil)
+	p, err := s.controller.ListPolicies(s.ctx, nil)
+	s.NoError(err)
+	s.NotNil(p)
+}
+
+func (s *preheatSuite) TestListPoliciesByProject() {
+	s.fakePolicyMgr.On("ListPoliciesByProject", s.ctx, int64(1), mock.Anything).Return([]*policy.Schema{}, nil)
+	p, err := s.controller.ListPoliciesByProject(s.ctx, 1, nil)
+	s.NoError(err)
+	s.NotNil(p)
+}