From fec09832f941f69c0219e2d7c82775749e2aab70 Mon Sep 17 00:00:00 2001
From: Rob Rudin <rob.rudin@marklogic.com>
Date: Wed, 22 May 2019 17:17:27 -0400
Subject: [PATCH] #362 Updating a task server now waits for a restart if
 necessary

---
 .../taskservers/UpdateTaskServerCommand.java  |  2 +-
 .../taskservers/TaskServerManager.java        | 20 ++++++++++++++++---
 .../taskservers/UpdateTaskServerTest.java     | 15 ++++++++------
 .../ml-config/task-servers/task-server.json   |  1 +
 4 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/src/main/java/com/marklogic/appdeployer/command/taskservers/UpdateTaskServerCommand.java b/src/main/java/com/marklogic/appdeployer/command/taskservers/UpdateTaskServerCommand.java
index 490d26c4..af6d8afc 100644
--- a/src/main/java/com/marklogic/appdeployer/command/taskservers/UpdateTaskServerCommand.java
+++ b/src/main/java/com/marklogic/appdeployer/command/taskservers/UpdateTaskServerCommand.java
@@ -37,7 +37,7 @@ public void execute(CommandContext context) {
 						logger.info("Processing file: " + f.getAbsolutePath());
 					}
 					String payload = copyFileToString(f, context);
-					mgr.updateTaskServer(taskServerName, payload);
+					mgr.updateTaskServer(taskServerName, payload, context.getAdminManager());
 				}
 			} else {
 				logResourceDirectoryNotFound(dir);
diff --git a/src/main/java/com/marklogic/mgmt/resource/taskservers/TaskServerManager.java b/src/main/java/com/marklogic/mgmt/resource/taskservers/TaskServerManager.java
index 5eb36b06..490d966c 100644
--- a/src/main/java/com/marklogic/mgmt/resource/taskservers/TaskServerManager.java
+++ b/src/main/java/com/marklogic/mgmt/resource/taskservers/TaskServerManager.java
@@ -2,7 +2,9 @@
 
 import com.marklogic.mgmt.AbstractManager;
 import com.marklogic.mgmt.ManageClient;
+import com.marklogic.mgmt.admin.AdminManager;
 import com.marklogic.rest.util.Fragment;
+import org.springframework.http.ResponseEntity;
 
 public class TaskServerManager extends AbstractManager {
 
@@ -12,12 +14,24 @@ public TaskServerManager(ManageClient manageClient) {
 		this.manageClient = manageClient;
 	}
 
-	public void updateTaskServer(String taskServerName, String payload) {
+	/**
+	 *
+	 * @param taskServerName
+	 * @param payload
+	 * @param adminManager required in the event that the update to the task server causes a restart
+	 */
+	public void updateTaskServer(String taskServerName, String payload, AdminManager adminManager) {
 		String path = format("/manage/v2/task-servers/%s/properties", taskServerName);
+
+		ResponseEntity<String> response;
 		if (payloadParser.isJsonPayload(payload)) {
-			manageClient.putJson(path, payload);
+			response = manageClient.putJson(path, payload);
 		} else {
-			manageClient.putXml(path, payload);
+			response = manageClient.putXml(path, payload);
+		}
+
+		if (response != null && response.getHeaders().getLocation() != null && adminManager != null) {
+			adminManager.waitForRestart();
 		}
 	}
 
diff --git a/src/test/java/com/marklogic/appdeployer/command/taskservers/UpdateTaskServerTest.java b/src/test/java/com/marklogic/appdeployer/command/taskservers/UpdateTaskServerTest.java
index d482dc79..8f892d58 100644
--- a/src/test/java/com/marklogic/appdeployer/command/taskservers/UpdateTaskServerTest.java
+++ b/src/test/java/com/marklogic/appdeployer/command/taskservers/UpdateTaskServerTest.java
@@ -1,7 +1,9 @@
 package com.marklogic.appdeployer.command.taskservers;
 
+import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.marklogic.appdeployer.AbstractAppDeployerTest;
 import com.marklogic.mgmt.resource.taskservers.TaskServerManager;
+import com.marklogic.mgmt.util.ObjectMapperFactory;
 import com.marklogic.rest.util.Fragment;
 import org.junit.Test;
 
@@ -10,6 +12,7 @@ public class UpdateTaskServerTest extends AbstractAppDeployerTest {
 	@Test
 	public void test() {
 		TaskServerManager mgr = new TaskServerManager(manageClient);
+		final int currentThreadCount = Integer.parseInt(mgr.getPropertiesAsXml().getElementValue("/node()/m:threads"));
 
 		initializeAppDeployer(new UpdateTaskServerCommand());
 		deploySampleApp();
@@ -20,13 +23,13 @@ public void test() {
 			assertEquals("false", xml.getElementValue("/m:task-server-properties/m:debug-allow"));
 			assertEquals("false", xml.getElementValue("/m:task-server-properties/m:profile-allow"));
 		} finally {
-			String payload = "{\n" +
-				"\t\"log-errors\": true,\n" +
-				"\t\"debug-allow\": true,\n" +
-				"\t\"profile-allow\": true\n" +
-				"}";
+			ObjectNode payload = ObjectMapperFactory.getObjectMapper().createObjectNode();
+			payload.put("threads", currentThreadCount);
+			payload.put("log-errors", true);
+			payload.put("debug-allow", true);
+			payload.put("profile-allow", true);
+			mgr.updateTaskServer("TaskServer", payload.toString(), adminManager);
 
-			mgr.updateTaskServer("TaskServer", payload);
 			Fragment xml = mgr.getPropertiesAsXml();
 			assertEquals("true", xml.getElementValue("/m:task-server-properties/m:log-errors"));
 			assertEquals("true", xml.getElementValue("/m:task-server-properties/m:debug-allow"));
diff --git a/src/test/resources/sample-app/src/main/ml-config/task-servers/task-server.json b/src/test/resources/sample-app/src/main/ml-config/task-servers/task-server.json
index ce57aa1c..fec86800 100644
--- a/src/test/resources/sample-app/src/main/ml-config/task-servers/task-server.json
+++ b/src/test/resources/sample-app/src/main/ml-config/task-servers/task-server.json
@@ -1,4 +1,5 @@
 {
+	"threads": 32,
 	"log-errors": false,
 	"debug-allow": false,
 	"profile-allow": false