Fix create team, update team missing units (#5188)

2 files changed, 28 insertions, 0 deletions

diff --git a/routers/api/v1/convert/convert.go b/routers/api/v1/convert/convert.go
index 1bfeae34bf..35416dea51 100644
--- a/routers/api/v1/convert/convert.go
+++ b/routers/api/v1/convert/convert.go
@@ -198,5 +198,6 @@ func ToTeam(team *models.Team) *api.Team {
 		Name:        team.Name,
 		Description: team.Description,
 		Permission:  team.Authorize.String(),
+		Units:       team.GetUnitNames(),
 	}
 }
diff --git a/routers/api/v1/org/team.go b/routers/api/v1/org/team.go
index 2e319a1831..8b67eda42f 100644
--- a/routers/api/v1/org/team.go
+++ b/routers/api/v1/org/team.go
@@ -90,6 +90,20 @@ func CreateTeam(ctx *context.APIContext, form api.CreateTeamOption) {
 		Description: form.Description,
 		Authorize:   models.ParseAccessMode(form.Permission),
 	}
+
+	unitTypes := models.FindUnitTypes(form.Units...)
+
+	if team.Authorize < models.AccessModeOwner {
+		var units = make([]*models.TeamUnit, 0, len(form.Units))
+		for _, tp := range unitTypes {
+			units = append(units, &models.TeamUnit{
+				OrgID: ctx.Org.Organization.ID,
+				Type:  tp,
+			})
+		}
+		team.Units = units
+	}
+
 	if err := models.NewTeam(team); err != nil {
 		if models.IsErrTeamAlreadyExist(err) {
 			ctx.Error(422, "", err)
@@ -128,6 +142,19 @@ func EditTeam(ctx *context.APIContext, form api.EditTeamOption) {
 	team.Name = form.Name
 	team.Description = form.Description
 	team.Authorize = models.ParseAccessMode(form.Permission)
+	unitTypes := models.FindUnitTypes(form.Units...)
+
+	if team.Authorize < models.AccessModeOwner {
+		var units = make([]*models.TeamUnit, 0, len(form.Units))
+		for _, tp := range unitTypes {
+			units = append(units, &models.TeamUnit{
+				OrgID: ctx.Org.Organization.ID,
+				Type:  tp,
+			})
+		}
+		team.Units = units
+	}
+
 	if err := models.UpdateTeam(team, true); err != nil {
 		ctx.Error(500, "EditTeam", err)
 		return