diff --git a/dolphinscheduler-api/src/main/java/org/apache/dolphinscheduler/api/service/impl/ProjectWorkerGroupRelationServiceImpl.java b/dolphinscheduler-api/src/main/java/org/apache/dolphinscheduler/api/service/impl/ProjectWorkerGroupRelationServiceImpl.java index 15cbcbb7fa..31dc113dbb 100644 --- a/dolphinscheduler-api/src/main/java/org/apache/dolphinscheduler/api/service/impl/ProjectWorkerGroupRelationServiceImpl.java +++ b/dolphinscheduler-api/src/main/java/org/apache/dolphinscheduler/api/service/impl/ProjectWorkerGroupRelationServiceImpl.java @@ -26,6 +26,7 @@ import org.apache.dolphinscheduler.common.constants.Constants; import org.apache.dolphinscheduler.dao.entity.Project; import org.apache.dolphinscheduler.dao.entity.ProjectWorkerGroup; import org.apache.dolphinscheduler.dao.entity.User; +import org.apache.dolphinscheduler.dao.entity.WorkerGroup; import org.apache.dolphinscheduler.dao.mapper.ProjectMapper; import org.apache.dolphinscheduler.dao.mapper.ProjectWorkerGroupMapper; import org.apache.dolphinscheduler.dao.mapper.ScheduleMapper; @@ -38,6 +39,7 @@ import org.apache.commons.lang3.StringUtils; import java.util.Date; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Objects; @@ -113,23 +115,25 @@ public class ProjectWorkerGroupRelationServiceImpl extends BaseServiceImpl } Set workerGroupNames = - workerGroupMapper.queryAllWorkerGroup().stream().map(item -> item.getName()).collect( + workerGroupMapper.queryAllWorkerGroup().stream().map(WorkerGroup::getName).collect( Collectors.toSet()); workerGroupNames.add(Constants.DEFAULT_WORKER_GROUP); - Set assignedWorkerGroupNames = workerGroups.stream().collect(Collectors.toSet()); + Set assignedWorkerGroupNames = new HashSet<>(workerGroups); Set difference = SetUtils.difference(assignedWorkerGroupNames, workerGroupNames); - if (difference.size() > 0) { + if (!difference.isEmpty()) { putMsg(result, Status.WORKER_GROUP_NOT_EXIST, difference.toString()); return result; } Set projectWorkerGroupNames = projectWorkerGroupMapper.selectList(new QueryWrapper() .lambda() - .eq(ProjectWorkerGroup::getProjectCode, projectCode)).stream().map(item -> item.getWorkerGroup()) + .eq(ProjectWorkerGroup::getProjectCode, projectCode)) + .stream() + .map(ProjectWorkerGroup::getWorkerGroup) .collect(Collectors.toSet()); difference = SetUtils.difference(projectWorkerGroupNames, assignedWorkerGroupNames); diff --git a/dolphinscheduler-api/src/test/java/org/apache/dolphinscheduler/api/service/ProjectWorkerGroupRelationServiceTest.java b/dolphinscheduler-api/src/test/java/org/apache/dolphinscheduler/api/service/ProjectWorkerGroupRelationServiceTest.java index a8e8b8beb2..6f6b7ca2d6 100644 --- a/dolphinscheduler-api/src/test/java/org/apache/dolphinscheduler/api/service/ProjectWorkerGroupRelationServiceTest.java +++ b/dolphinscheduler-api/src/test/java/org/apache/dolphinscheduler/api/service/ProjectWorkerGroupRelationServiceTest.java @@ -17,13 +17,17 @@ package org.apache.dolphinscheduler.api.service; +import static org.apache.dolphinscheduler.api.utils.ServiceTestUtil.getAdminUser; +import static org.apache.dolphinscheduler.api.utils.ServiceTestUtil.getGeneralUser; + +import org.apache.dolphinscheduler.api.AssertionsHelper; import org.apache.dolphinscheduler.api.enums.Status; import org.apache.dolphinscheduler.api.service.impl.ProjectWorkerGroupRelationServiceImpl; import org.apache.dolphinscheduler.api.utils.Result; import org.apache.dolphinscheduler.common.constants.Constants; -import org.apache.dolphinscheduler.common.enums.UserType; import org.apache.dolphinscheduler.dao.entity.Project; import org.apache.dolphinscheduler.dao.entity.ProjectWorkerGroup; +import org.apache.dolphinscheduler.dao.entity.TaskDefinition; import org.apache.dolphinscheduler.dao.entity.User; import org.apache.dolphinscheduler.dao.entity.WorkerGroup; import org.apache.dolphinscheduler.dao.mapper.ProjectMapper; @@ -33,6 +37,7 @@ import org.apache.dolphinscheduler.dao.mapper.TaskDefinitionMapper; import org.apache.dolphinscheduler.dao.mapper.WorkerGroupMapper; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Map; @@ -77,27 +82,87 @@ public class ProjectWorkerGroupRelationServiceTest { @Test public void testAssignWorkerGroupsToProject() { + User generalUser = getGeneralUser(); User loginUser = getAdminUser(); + // no permission + Result result = projectWorkerGroupRelationService.assignWorkerGroupsToProject(generalUser, projectCode, + getWorkerGroups()); + Assertions.assertEquals(Status.USER_NO_OPERATION_PERM.getCode(), result.getCode()); + + // project code is null + result = projectWorkerGroupRelationService.assignWorkerGroupsToProject(loginUser, null, + getWorkerGroups()); + Assertions.assertEquals(Status.PROJECT_NOT_EXIST.getCode(), result.getCode()); + + // worker group is empty + result = projectWorkerGroupRelationService.assignWorkerGroupsToProject(loginUser, projectCode, + Collections.emptyList()); + Assertions.assertEquals(Status.WORKER_GROUP_TO_PROJECT_IS_EMPTY.getCode(), result.getCode()); + + // project not exists Mockito.when(projectMapper.queryByCode(projectCode)).thenReturn(null); - Result result = projectWorkerGroupRelationService.assignWorkerGroupsToProject(loginUser, projectCode, + result = projectWorkerGroupRelationService.assignWorkerGroupsToProject(loginUser, projectCode, getWorkerGroups()); Assertions.assertEquals(Status.PROJECT_NOT_EXIST.getCode(), result.getCode()); + // worker group not exists WorkerGroup workerGroup = new WorkerGroup(); workerGroup.setName("test"); Mockito.when(projectMapper.queryByCode(Mockito.anyLong())).thenReturn(getProject()); - Mockito.when(workerGroupMapper.queryAllWorkerGroup()).thenReturn(Lists.newArrayList(workerGroup)); + Mockito.when(workerGroupMapper.queryAllWorkerGroup()).thenReturn(Collections.singletonList(workerGroup)); + result = projectWorkerGroupRelationService.assignWorkerGroupsToProject(loginUser, projectCode, + getDiffWorkerGroups()); + Assertions.assertEquals(Status.WORKER_GROUP_NOT_EXIST.getCode(), result.getCode()); + + // db insertion fail + Mockito.when(workerGroupMapper.queryAllWorkerGroup()).thenReturn(Collections.singletonList(workerGroup)); + Mockito.when(projectWorkerGroupMapper.insert(Mockito.any())).thenReturn(-1); + AssertionsHelper.assertThrowsServiceException(Status.ASSIGN_WORKER_GROUP_TO_PROJECT_ERROR, + () -> projectWorkerGroupRelationService.assignWorkerGroupsToProject(loginUser, projectCode, + getWorkerGroups())); + + // success Mockito.when(projectWorkerGroupMapper.insert(Mockito.any())).thenReturn(1); result = projectWorkerGroupRelationService.assignWorkerGroupsToProject(loginUser, projectCode, getWorkerGroups()); Assertions.assertEquals(Status.SUCCESS.getCode(), result.getCode()); + + // success when there is diff between current wg and assigned wg + Mockito.when(projectWorkerGroupMapper.selectList(Mockito.any())) + .thenReturn(Collections.singletonList(getDiffProjectWorkerGroup())); + Mockito.when(projectWorkerGroupMapper.delete(Mockito.any())).thenReturn(1); + result = projectWorkerGroupRelationService.assignWorkerGroupsToProject(loginUser, projectCode, + getWorkerGroups()); + Assertions.assertEquals(Status.SUCCESS.getCode(), result.getCode()); + + // db deletion fail + Mockito.when(projectWorkerGroupMapper.delete(Mockito.any())).thenReturn(-1); + AssertionsHelper.assertThrowsServiceException(Status.ASSIGN_WORKER_GROUP_TO_PROJECT_ERROR, + () -> projectWorkerGroupRelationService.assignWorkerGroupsToProject(loginUser, projectCode, + getWorkerGroups())); + + // fail when wg is referenced by task definition + Mockito.when(taskDefinitionMapper.queryAllDefinitionList(Mockito.anyLong())) + .thenReturn(Collections.singletonList(getTaskDefinitionWithDiffWorkerGroup())); + AssertionsHelper.assertThrowsServiceException(Status.USED_WORKER_GROUP_EXISTS, + () -> projectWorkerGroupRelationService.assignWorkerGroupsToProject(loginUser, projectCode, + getWorkerGroups())); } @Test public void testQueryWorkerGroupsByProject() { + // no permission + Mockito.when(projectService.hasProjectAndPerm(Mockito.any(), Mockito.any(), Mockito.anyMap(), Mockito.any())) + .thenReturn(false); + Map result = + projectWorkerGroupRelationService.queryWorkerGroupsByProject(getGeneralUser(), projectCode); + + Assertions.assertTrue(result.isEmpty()); + + // success Mockito.when(projectService.hasProjectAndPerm(Mockito.any(), Mockito.any(), Mockito.anyMap(), Mockito.any())) .thenReturn(true); @@ -113,8 +178,7 @@ public class ProjectWorkerGroupRelationServiceTest { Mockito.when(scheduleMapper.querySchedulerListByProjectName(Mockito.any())) .thenReturn(Lists.newArrayList()); - Map result = - projectWorkerGroupRelationService.queryWorkerGroupsByProject(getGeneralUser(), projectCode); + result = projectWorkerGroupRelationService.queryWorkerGroupsByProject(getGeneralUser(), projectCode); ProjectWorkerGroup[] actualValue = ((List) result.get(Constants.DATA_LIST)).toArray(new ProjectWorkerGroup[0]); @@ -126,20 +190,8 @@ public class ProjectWorkerGroupRelationServiceTest { return Lists.newArrayList("default"); } - private User getGeneralUser() { - User loginUser = new User(); - loginUser.setUserType(UserType.GENERAL_USER); - loginUser.setUserName("userName"); - loginUser.setId(1); - return loginUser; - } - - private User getAdminUser() { - User loginUser = new User(); - loginUser.setUserType(UserType.ADMIN_USER); - loginUser.setUserName("userName"); - loginUser.setId(1); - return loginUser; + private List getDiffWorkerGroups() { + return Lists.newArrayList("default", "new"); } private Project getProject() { @@ -158,4 +210,20 @@ public class ProjectWorkerGroupRelationServiceTest { projectWorkerGroup.setWorkerGroup("default"); return projectWorkerGroup; } + + private ProjectWorkerGroup getDiffProjectWorkerGroup() { + ProjectWorkerGroup projectWorkerGroup = new ProjectWorkerGroup(); + projectWorkerGroup.setId(2); + projectWorkerGroup.setProjectCode(projectCode); + projectWorkerGroup.setWorkerGroup("new"); + return projectWorkerGroup; + } + + private TaskDefinition getTaskDefinitionWithDiffWorkerGroup() { + TaskDefinition taskDefinition = new TaskDefinition(); + taskDefinition.setProjectCode(projectCode); + taskDefinition.setId(1); + taskDefinition.setWorkerGroup("new"); + return taskDefinition; + } }