[Fix-17604][API] Correct logic for assigning and remove worker groups to project

dolphinscheduler-api/src/main/java/org/apache/dolphinscheduler/api/service/impl/ProjectWorkerGroupRelationServiceImpl.java

@@ -125,26 +125,32 @@ public Result assignWorkerGroupsToProject(User loginUser, Long projectCode, List
         // check if assign worker group exists in the project
         Set<String> projectWorkerGroupNames =
                 projectWorkerGroupDao.queryAssignedWorkerGroupNamesByProjectCode(projectCode);
-        difference = SetUtils.difference(unauthorizedWorkerGroupNames, projectWorkerGroupNames);
+        Set<String> needDeletedWorkerGroups =
+                SetUtils.difference(projectWorkerGroupNames, unauthorizedWorkerGroupNames);
         Date now = new Date();
-        if (CollectionUtils.isNotEmpty(difference)) {
-            Set<String> usedWorkerGroups = getAllUsedWorkerGroups(project);
 
-            if (CollectionUtils.isNotEmpty(usedWorkerGroups) && usedWorkerGroups.containsAll(difference)) {
+        if (CollectionUtils.isNotEmpty(needDeletedWorkerGroups)) {
+            Set<String> usedWorkerGroups = getAllUsedWorkerGroups(project);
+            if (CollectionUtils.isNotEmpty(usedWorkerGroups) && usedWorkerGroups.containsAll(needDeletedWorkerGroups)) {
                 throw new ServiceException(Status.USED_WORKER_GROUP_EXISTS,
-                        SetUtils.intersection(usedWorkerGroups, difference).toSet());
+                        SetUtils.intersection(usedWorkerGroups, needDeletedWorkerGroups).toSet());
             }
-
             boolean deleted =
-                    projectWorkerGroupDao.deleteByProjectCodeAndWorkerGroups(projectCode, new ArrayList<>(difference));
+                    projectWorkerGroupDao.deleteByProjectCodeAndWorkerGroups(projectCode,
+                            new ArrayList<>(needDeletedWorkerGroups));
             if (deleted) {
-                log.info("Success to delete worker groups [{}] for the project [{}] .", difference, project.getName());
+                log.info("Success to delete worker groups [{}] for the project [{}] .", needDeletedWorkerGroups,
+                        project.getName());
             } else {
-                log.error("Failed to delete worker groups [{}] for the project [{}].", difference, project.getName());
+                log.error("Failed to delete worker groups [{}] for the project [{}].", needDeletedWorkerGroups,
+                        project.getName());
                 throw new ServiceException(Status.ASSIGN_WORKER_GROUP_TO_PROJECT_ERROR);
             }
-
-            difference.forEach(workerGroupName -> {
+        }
+        Set<String> needAssignedWorkerGroups =
+                SetUtils.difference(unauthorizedWorkerGroupNames, projectWorkerGroupNames);
+        if (CollectionUtils.isNotEmpty(needAssignedWorkerGroups)) {
+            needAssignedWorkerGroups.forEach(workerGroupName -> {
                 ProjectWorkerGroup projectWorkerGroup = new ProjectWorkerGroup();
                 projectWorkerGroup.setProjectCode(projectCode);
                 projectWorkerGroup.setWorkerGroup(workerGroupName);

dolphinscheduler-api/src/test/java/org/apache/dolphinscheduler/api/service/ProjectWorkerGroupRelationServiceTest.java

@@ -146,11 +146,14 @@ public void testAssignWorkerGroupsToProject() {
                         getWorkerGroups()));
 
         // fail when wg is referenced by task definition
+        // test case: project all wg: test, task used wg: test, new wg: test1
         Mockito.when(taskDefinitionDao.queryAllTaskDefinitionWorkerGroups(Mockito.anyLong()))
                 .thenReturn(Collections.singletonList(getProjectWorkerGroup().getWorkerGroup()));
+        Mockito.when(projectWorkerGroupDao.queryAssignedWorkerGroupNamesByProjectCode(Mockito.any()))
+                .thenReturn(Sets.newHashSet(getProjectWorkerGroup().getWorkerGroup()));
         AssertionsHelper.assertThrowsServiceException(Status.USED_WORKER_GROUP_EXISTS,
                 () -> projectWorkerGroupRelationService.assignWorkerGroupsToProject(loginUser, projectCode,
-                        getWorkerGroups()));
+                        getUnusedWorkerGroups()));
     }
 
     @Test
@@ -192,6 +195,10 @@ private List<String> getWorkerGroups() {
         return Lists.newArrayList("test");
     }
 
+    private List<String> getUnusedWorkerGroups() {
+        return Lists.newArrayList("test1");
+    }
+
     private List<String> getDiffWorkerGroups() {
         return Lists.newArrayList("test3", "new");
     }

dolphinscheduler-e2e/dolphinscheduler-e2e-case/src/test/java/org/apache/dolphinscheduler/e2e/cases/ProjectE2ETest.java

@@ -22,6 +22,8 @@
 import org.apache.dolphinscheduler.e2e.core.DolphinScheduler;
 import org.apache.dolphinscheduler.e2e.pages.LoginPage;
 import org.apache.dolphinscheduler.e2e.pages.project.ProjectPage;
+import org.apache.dolphinscheduler.e2e.pages.security.SecurityPage;
+import org.apache.dolphinscheduler.e2e.pages.security.WorkerGroupPage;
 
 import java.util.UUID;
 
@@ -39,6 +41,7 @@ class ProjectE2ETest {
     private static final String project = "test-project-" + UUID.randomUUID();
 
     private static final String workerGroup = "default";
+    private static final String workerGroupTest = "test-wg-" + UUID.randomUUID();
 
     private static RemoteWebDriver browser;
 
@@ -71,6 +74,26 @@ void testAssignWorkerGroup() {
         page.verifyAssignedWorkerGroup(project, workerGroup);
     }
 
+    @Test
+    @Order(6)
+    void testAssignNewWorkerGroup() {
+        WorkerGroupPage workerGroupPage =
+                new WorkerGroupPage(browser).goToNav(SecurityPage.class).goToTab(WorkerGroupPage.class);
+        workerGroupPage.create(workerGroupTest);
+        workerGroupPage.verifyWorkerGroupCreated(workerGroupTest);
+        ProjectPage projectPage = new ProjectPage(browser).goToNav(ProjectPage.class);
+        projectPage.assignWorkerGroup(project, workerGroupTest);
+        projectPage.verifyAssignedWorkerGroup(project, workerGroupTest);
+    }
+
+    @Test
+    @Order(7)
+    void testRemoveWorkerGroup() {
+        final ProjectPage page = new ProjectPage(browser);
+        page.removeWorkerGroup(project, workerGroup);
+        page.verifyRemovedWorkerGroup(project, workerGroup);
+    }
+
     @Test
     @Order(30)
     void testDeleteProject() {

dolphinscheduler-e2e/dolphinscheduler-e2e-case/src/test/java/org/apache/dolphinscheduler/e2e/pages/project/ProjectPage.java

@@ -130,7 +130,49 @@ public ProjectPage verifyAssignedWorkerGroup(String project, String workerGroup)
 
         return this;
     }
+    public ProjectPage removeWorkerGroup(String project, String workerGroup) {
+        projectList()
+                .stream()
+                .filter(it -> it.getText().contains(project))
+                .findFirst()
+                .orElseThrow(() -> new RuntimeException("Can not find project: " + project))
+                .findElement(By.className("assign-worker-group-btn")).click();
+
+        WebElement workerGroupItem = assignWorkerGroupForm.sourceWorkerGroupItems()
+                .stream()
+                .filter(it -> it.findElement(By.className("n-transfer-list-item__label"))
+                        .getText().contains(workerGroup))
+                .filter(it -> {
+                    WebElement checkbox = it.findElement(By.className("n-checkbox"));
+                    return "true".equals(checkbox.getAttribute("aria-checked"));
+                })
+                .findFirst()
+                .orElseThrow(() -> new RuntimeException("Can not find selected source worker group: " + workerGroup));
+
+        workerGroupItem.click();
+        assignWorkerGroupForm.buttonSubmit().click();
+        return this;
+    }
+    public ProjectPage verifyRemovedWorkerGroup(String project, String workerGroup) {
+        projectList()
+                .stream()
+                .filter(it -> it.getText().contains(project))
+                .findFirst()
+                .orElseThrow(() -> new RuntimeException("Can not find project: " + project))
+                .findElement(By.className("assign-worker-group-btn")).click();
+
+        assignWorkerGroupForm.targetWorkerGroups()
+                .stream()
+                .filter(it -> it.getText().contains(workerGroup))
+                .findFirst()
+                .ifPresent(it -> {
+                    throw new RuntimeException(
+                            "Worker group should have been deleted but still exists: " + workerGroup);
+                });
+        assignWorkerGroupForm.buttonCancel().click();
 
+        return this;
+    }
     public ProjectDetailPage goTo(String project) {
         projectList().stream()
                 .filter(it -> it.getText().contains(project))
@@ -180,6 +222,13 @@ public class AssignWorkerGroupForm {
         })
         private List<WebElement> targetWorkerGroups;
 
+        @FindBys({
+                @FindBy(className = "assign-worker-group-modal"),
+                @FindBy(className = "n-transfer-list--source"),
+                @FindBy(className = "n-transfer-list-item--source")
+        })
+        private List<WebElement> sourceWorkerGroupItems;
+
         @FindBys({
                 @FindBy(className = "assign-worker-group-modal"),
                 @FindBy(className = "btn-submit"),

dolphinscheduler-e2e/dolphinscheduler-e2e-case/src/test/java/org/apache/dolphinscheduler/e2e/pages/security/WorkerGroupPage.java

@@ -17,8 +17,12 @@
 
 package org.apache.dolphinscheduler.e2e.pages.security;
 
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.testcontainers.shaded.org.awaitility.Awaitility.await;
+
 import org.apache.dolphinscheduler.e2e.pages.common.NavBarPage;
 
+import java.time.Duration;
 import java.util.List;
 
 import lombok.Getter;
@@ -99,6 +103,12 @@ public WorkerGroupPage delete(String Worker) {
 
         return this;
     }
+    public void verifyWorkerGroupCreated(String workerGroupName) {
+        await().atMost(Duration.ofMinutes(1)).untilAsserted(() -> assertThat(workerGroupList())
+                .as("workerGroup list should contain newly-created workerGroup")
+                .extracting(WebElement::getText)
+                .anyMatch(it -> it.contains(workerGroupName)));
+    }
 
     @Getter
     public class WorkerGroupForm {