Create file on shutdown to indicate that data prepper is shutdown

Signed-off-by: Taylor Gray <[email protected]>
opensearch-project · Feb 26, 2024 · 361c6d5 · 361c6d5
1 parent ea8d5f5
commit 361c6d5
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 2 deletions.
diff --git a/...repper-core/src/main/java/org/opensearch/dataprepper/pipeline/server/ShutdownHandler.java b/...repper-core/src/main/java/org/opensearch/dataprepper/pipeline/server/ShutdownHandler.java
@@ -14,6 +14,9 @@
 import javax.ws.rs.HttpMethod;
 import java.io.IOException;
 import java.net.HttpURLConnection;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
 
 /**
  * HttpHandler to handle requests to shut down the data prepper instance
@@ -22,6 +25,9 @@ public class ShutdownHandler implements HttpHandler {
     private final DataPrepper dataPrepper;
     private static final Logger LOG = LoggerFactory.getLogger(ShutdownHandler.class);
 
+    static final Path SHUTDOWN_FILE_PATH = Path.of("Is-Shutdown.log");
+    static final String SHUTDOWN_MESSAGE = "Data Prepper is shut down.";
+
     public ShutdownHandler(final DataPrepper dataPrepper) {
         this.dataPrepper = dataPrepper;
     }
@@ -44,6 +50,7 @@ public void handle(final HttpExchange exchange) throws IOException {
             exchange.sendResponseHeaders(HttpURLConnection.HTTP_INTERNAL_ERROR, 0);
         } finally {
             exchange.getResponseBody().close();
+            Files.write(SHUTDOWN_FILE_PATH, SHUTDOWN_MESSAGE.getBytes(StandardCharsets.UTF_8));
             dataPrepper.shutdownServers();
         }
     }

diff --git a/...er-core/src/test/java/org/opensearch/dataprepper/pipeline/server/ShutdownHandlerTest.java b/...er-core/src/test/java/org/opensearch/dataprepper/pipeline/server/ShutdownHandlerTest.java
@@ -13,19 +13,25 @@
 import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.InjectMocks;
 import org.mockito.Mock;
+import org.mockito.MockedStatic;
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.opensearch.dataprepper.DataPrepper;
 
 import javax.ws.rs.HttpMethod;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.net.HttpURLConnection;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
 
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mockStatic;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.opensearch.dataprepper.pipeline.server.ShutdownHandler.SHUTDOWN_FILE_PATH;
+import static org.opensearch.dataprepper.pipeline.server.ShutdownHandler.SHUTDOWN_MESSAGE;
 
 @ExtendWith(MockitoExtension.class)
 public class ShutdownHandlerTest {
@@ -52,7 +58,13 @@ public void testWhenShutdownWithPostRequestThenResponseWritten() throws IOExcept
         when(exchange.getRequestMethod())
                 .thenReturn(HttpMethod.POST);
 
-        shutdownHandler.handle(exchange);
+        try (final MockedStatic<Files> filesMockedStatic = mockStatic(Files.class)) {
+            filesMockedStatic.when(() -> Files.write(SHUTDOWN_FILE_PATH, SHUTDOWN_MESSAGE.getBytes(StandardCharsets.UTF_8)))
+                    .thenReturn(SHUTDOWN_FILE_PATH);
+            shutdownHandler.handle(exchange);
+
+            filesMockedStatic.verify(() -> Files.write(SHUTDOWN_FILE_PATH, SHUTDOWN_MESSAGE.getBytes(StandardCharsets.UTF_8)), times(1));
+        }
 
         verify(dataPrepper, times(1))
                 .shutdownPipelines();
@@ -62,6 +74,8 @@ public void testWhenShutdownWithPostRequestThenResponseWritten() throws IOExcept
                 .close();
         verify(dataPrepper, times(1))
                 .shutdownServers();
+
+
     }
 
     @ParameterizedTest
@@ -84,11 +98,20 @@ public void testHandleException() throws IOException {
                 .thenReturn(HttpMethod.POST);
         doThrow(RuntimeException.class).when(dataPrepper).shutdownPipelines();
 
-        shutdownHandler.handle(exchange);
+        try (final MockedStatic<Files> filesMockedStatic = mockStatic(Files.class)) {
+            filesMockedStatic.when(() -> Files.write(SHUTDOWN_FILE_PATH, SHUTDOWN_MESSAGE.getBytes(StandardCharsets.UTF_8)))
+                    .thenReturn(SHUTDOWN_FILE_PATH);
+            shutdownHandler.handle(exchange);
+
+            filesMockedStatic.verify(() -> Files.write(SHUTDOWN_FILE_PATH, SHUTDOWN_MESSAGE.getBytes(StandardCharsets.UTF_8)), times(1));
+        }
 
         verify(exchange, times(1))
                 .sendResponseHeaders(eq(HttpURLConnection.HTTP_INTERNAL_ERROR), eq(0L));
         verify(responseBody, times(1))
                 .close();
+
+        verify(dataPrepper, times(1))
+                .shutdownServers();
     }
 }