From c6ae1388eca730016323d99353dd03b1ddc69427 Mon Sep 17 00:00:00 2001
From: wang yan <wangyan@vmware.com>
Date: Thu, 10 Jan 2019 14:50:35 +0800
Subject: [PATCH] Add dao of robot account

Signed-off-by: wang yan <wangyan@vmware.com>
---
 src/common/dao/robot.go      |  86 ++++++++++++++++++++++++++++
 src/common/dao/robot_test.go | 107 +++++++++++++++++++++++++++++++++++
 2 files changed, 193 insertions(+)
 create mode 100644 src/common/dao/robot.go
 create mode 100644 src/common/dao/robot_test.go

diff --git a/src/common/dao/robot.go b/src/common/dao/robot.go
new file mode 100644
index 000000000..54c760374
--- /dev/null
+++ b/src/common/dao/robot.go
@@ -0,0 +1,86 @@
+// Copyright Project Harbor Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dao
+
+import (
+	"fmt"
+	"github.com/astaxie/beego/orm"
+	"github.com/goharbor/harbor/src/common/models"
+	"time"
+)
+
+// AddRobot ...
+func AddRobot(robot *models.Robot) (int64, error) {
+	now := time.Now()
+	robot.CreationTime = now
+	robot.UpdateTime = now
+	return GetOrmer().Insert(robot)
+}
+
+// GetRobotByID ...
+func GetRobotByID(id int64) (*models.Robot, error) {
+	robot := &models.Robot{
+		ID: id,
+	}
+	if err := GetOrmer().Read(robot); err != nil {
+		if err == orm.ErrNoRows {
+			return nil, nil
+		}
+		return nil, err
+	}
+
+	return robot, nil
+}
+
+// ListRobots list robots according to the query conditions
+func ListRobots(query *models.RobotQuery) ([]*models.Robot, error) {
+	qs := getRobotQuerySetter(query)
+	if query.Size > 0 {
+		qs = qs.Limit(query.Size)
+		if query.Page > 0 {
+			qs = qs.Offset((query.Page - 1) * query.Size)
+		}
+	}
+	qs = qs.OrderBy("Name")
+
+	robots := []*models.Robot{}
+	_, err := qs.All(&robots)
+	return robots, err
+}
+
+func getRobotQuerySetter(query *models.RobotQuery) orm.QuerySeter {
+	qs := GetOrmer().QueryTable(&models.Robot{})
+	if len(query.Name) > 0 {
+		qs = qs.Filter("Name", query.Name)
+	}
+	if query.ProjectID != 0 {
+		qs = qs.Filter("ProjectID", query.ProjectID)
+	}
+	qs = qs.Filter("Disabled", false)
+	return qs
+}
+
+// DisableRobot ...
+func DisableRobot(id int64) error {
+	robot, err := GetRobotByID(id)
+	if err != nil {
+		return err
+	}
+	robot.Name = fmt.Sprintf("%s#%d", robot.Name, robot.ID)
+	robot.UpdateTime = time.Now()
+	robot.Disabled = true
+	_, err = GetOrmer().Update(robot, "Name", "UpdateTime", "Disabled")
+	return err
+}
diff --git a/src/common/dao/robot_test.go b/src/common/dao/robot_test.go
new file mode 100644
index 000000000..b047018fd
--- /dev/null
+++ b/src/common/dao/robot_test.go
@@ -0,0 +1,107 @@
+// Copyright Project Harbor Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dao
+
+import (
+	"testing"
+
+	"github.com/goharbor/harbor/src/common/models"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestAddRobot(t *testing.T) {
+	robotName := "test1"
+	robot := &models.Robot{
+		Name:        robotName,
+		Token:       "rKgjKEMpMEK23zqejkWn5GIVvgJps1vKACTa6tnGXXyOlOTsXFESccDvgaJx047q1",
+		Description: "test1 description",
+		ProjectID:   1,
+	}
+
+	// add
+	id, err := AddRobot(robot)
+	require.Nil(t, err)
+	robot.ID = id
+
+	require.Nil(t, err)
+	assert.NotNil(t, id)
+
+}
+
+func TestGetRobot(t *testing.T) {
+	robotName := "test2"
+	robot := &models.Robot{
+		Name:        robotName,
+		Token:       "rKgjKEMpMEK23zqejkWn5GIVvgJps1vKACTa6tnGXXyOlOTsXFESccDvgaJx047q2",
+		Description: "test2 description",
+		ProjectID:   1,
+	}
+
+	// add
+	id, err := AddRobot(robot)
+	require.Nil(t, err)
+	robot.ID = id
+
+	robot, err = GetRobotByID(id)
+	require.Nil(t, err)
+	assert.Equal(t, robotName, robot.Name)
+
+}
+
+func TestListRobots(t *testing.T) {
+	robotName := "test3"
+	robot := &models.Robot{
+		Name:        robotName,
+		Token:       "rKgjKEMpMEK23zqejkWn5GIVvgJps1vKACTa6tnGXXyOlOTsXFESccDvgaJx047q3",
+		Description: "test3 description",
+		ProjectID:   1,
+	}
+
+	_, err := AddRobot(robot)
+	require.Nil(t, err)
+
+	robots, err := ListRobots(&models.RobotQuery{
+		ProjectID: 1,
+	})
+	require.Nil(t, err)
+	assert.Equal(t, 3, len(robots))
+
+}
+
+func TestDisableRobot(t *testing.T) {
+	robotName := "test4"
+	robot := &models.Robot{
+		Name:        robotName,
+		Token:       "rKgjKEMpMEK23zqejkWn5GIVvgJps1vKACTa6tnGXXyOlOTsXFESccDvgaJx047q4",
+		Description: "test4 description",
+		ProjectID:   1,
+	}
+
+	// add
+	id, err := AddRobot(robot)
+	require.Nil(t, err)
+
+	// Disable
+	err = DisableRobot(id)
+	require.Nil(t, err)
+
+	// Get
+	robot, err = GetRobotByID(id)
+	require.Nil(t, err)
+	assert.Equal(t, true, robot.Disabled)
+
+}