From ec6dd2a720e39add0fdabfe7898259af63268671 Mon Sep 17 00:00:00 2001 From: Rishu Yadav Date: Mon, 28 Oct 2024 12:57:08 +0530 Subject: [PATCH 1/2] feat: execution name sanitization and unit test for stepfn --- .../emodb/queue/api/DedupQueueService.java | 5 + .../emodb/queue/api/QueueService.java | 5 +- .../queue/core/AbstractQueueService.java | 5 +- .../core/stepfn/StepFunctionService.java | 34 +++ .../core/stepfn/StepFunctionServiceTest.java | 199 ++++++++++++++++++ 5 files changed, 243 insertions(+), 5 deletions(-) create mode 100644 queue/src/test/java/com/bazaarvoice/emodb/queue/core/stepfn/StepFunctionServiceTest.java diff --git a/queue-api/src/main/java/com/bazaarvoice/emodb/queue/api/DedupQueueService.java b/queue-api/src/main/java/com/bazaarvoice/emodb/queue/api/DedupQueueService.java index a6dc77515..5952c87b0 100644 --- a/queue-api/src/main/java/com/bazaarvoice/emodb/queue/api/DedupQueueService.java +++ b/queue-api/src/main/java/com/bazaarvoice/emodb/queue/api/DedupQueueService.java @@ -15,6 +15,7 @@ public interface DedupQueueService extends BaseQueueService { void sendAll(Map> messagesByQueue); + //Overloaded sendAll method to send to cassandra void sendAll(String queue, Collectionmessages, boolean isFlush); /** @@ -26,6 +27,10 @@ public interface DedupQueueService extends BaseQueueService { */ long getMessageCount(String queue); + default long getUncachedSize(String queue){ + return 0; + } + /** * Counts the total number of messages for the specified queue, accurate up to the specified limit. Beyond the * specified limit the message count will be a rough estimate, allowing the caller to make the trade-off between diff --git a/queue-api/src/main/java/com/bazaarvoice/emodb/queue/api/QueueService.java b/queue-api/src/main/java/com/bazaarvoice/emodb/queue/api/QueueService.java index 4d533c755..14b2bf215 100644 --- a/queue-api/src/main/java/com/bazaarvoice/emodb/queue/api/QueueService.java +++ b/queue-api/src/main/java/com/bazaarvoice/emodb/queue/api/QueueService.java @@ -13,7 +13,6 @@ public interface QueueService extends BaseQueueService { void sendAll(String queue, Collection messages); - void sendAll(Map> messagesByQueue); //Overloaded sendAll method to send to cassandra @@ -28,6 +27,10 @@ public interface QueueService extends BaseQueueService { */ long getMessageCount(String queue); + default long getUncachedSize(String queue){ + return 0; + } + /** * Counts the total number of messages for the specified queue, accurate up to the specified limit. Beyond the * specified limit the message count will be a rough estimate, allowing the caller to make the trade-off between diff --git a/queue/src/main/java/com/bazaarvoice/emodb/queue/core/AbstractQueueService.java b/queue/src/main/java/com/bazaarvoice/emodb/queue/core/AbstractQueueService.java index 0793a08b2..1a6589f84 100644 --- a/queue/src/main/java/com/bazaarvoice/emodb/queue/core/AbstractQueueService.java +++ b/queue/src/main/java/com/bazaarvoice/emodb/queue/core/AbstractQueueService.java @@ -454,12 +454,9 @@ private void startStepFunctionExecution(Map parameters, String q String inputPayload = createInputPayload(queueThreshold, batchSize, queueType, queueName, topic, interval); - // Create the timestamp - String timestamp = String.valueOf(System.currentTimeMillis()); // Current time in milliseconds // Check if queueType is "dedupq" and prepend "D" to execution name if true - String executionName = (queueType.equalsIgnoreCase("dedupq") ? "D_" : "") + queueName + "_" + timestamp; - + String executionName = (queueType.equalsIgnoreCase("dedupq") ? "D_" : "") + queueName ; // Start the Step Function execution stepFunctionService.startExecution(stateMachineArn, inputPayload, executionName); diff --git a/queue/src/main/java/com/bazaarvoice/emodb/queue/core/stepfn/StepFunctionService.java b/queue/src/main/java/com/bazaarvoice/emodb/queue/core/stepfn/StepFunctionService.java index ec4165d03..6a9597884 100644 --- a/queue/src/main/java/com/bazaarvoice/emodb/queue/core/stepfn/StepFunctionService.java +++ b/queue/src/main/java/com/bazaarvoice/emodb/queue/core/stepfn/StepFunctionService.java @@ -7,6 +7,7 @@ import com.amazonaws.services.stepfunctions.model.StartExecutionResult; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.time.Instant; /** * Service to interact with AWS Step Functions using AWS SDK v1. @@ -42,6 +43,10 @@ public void startExecution(String stateMachineArn, String inputPayload, String e logger.warn("Input payload is null; using empty JSON object"); inputPayload = "{}"; // Default to empty payload if null } + // Create the timestamp + String timestamp = String.valueOf(Instant.now().getEpochSecond()); + // Append the timestamp to the initial execution name + executionName = sanitizeExecutionName(executionName) + "_" + timestamp; try { StartExecutionRequest startExecutionRequest = new StartExecutionRequest() @@ -59,4 +64,33 @@ public void startExecution(String stateMachineArn, String inputPayload, String e throw e; } } + + /** + * Sanitizes the execution name by replacing invalid characters with underscores + * and truncating if needed. + */ + public String sanitizeExecutionName(String executionName) { + if (executionName == null || executionName.isEmpty()) { + throw new IllegalArgumentException("Execution name cannot be null or empty"); + } + executionName = executionName.trim(); + // Replace invalid characters with underscores + String sanitized = executionName.replaceAll("[^a-zA-Z0-9\\-_]", "_"); + + // Check if the sanitized name is empty or consists only of underscores + if (sanitized.isEmpty() || sanitized.replaceAll("_", "").isEmpty()) { + throw new IllegalArgumentException("Execution name cannot contain only invalid characters"); + } + + // Truncate from the beginning if length exceeds 69 characters + if (sanitized.length() > 69) { + sanitized = sanitized.substring(sanitized.length() - 69); + } + + // Log the updated execution name if it has changed + if (!sanitized.equals(executionName)) { + logger.info("Updated execution name: {}", sanitized); + } + return sanitized; + } } \ No newline at end of file diff --git a/queue/src/test/java/com/bazaarvoice/emodb/queue/core/stepfn/StepFunctionServiceTest.java b/queue/src/test/java/com/bazaarvoice/emodb/queue/core/stepfn/StepFunctionServiceTest.java new file mode 100644 index 000000000..09d5fecbe --- /dev/null +++ b/queue/src/test/java/com/bazaarvoice/emodb/queue/core/stepfn/StepFunctionServiceTest.java @@ -0,0 +1,199 @@ +package com.bazaarvoice.emodb.queue.core.stepfn; + +import com.amazonaws.services.stepfunctions.AWSStepFunctions; +import com.amazonaws.services.stepfunctions.model.StartExecutionRequest; +import com.amazonaws.services.stepfunctions.model.StartExecutionResult; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import java.lang.reflect.Field; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.*; +import static org.testng.Assert.*; + +public class StepFunctionServiceTest { + + private StepFunctionService stepFunctionService; + + @Mock + private AWSStepFunctions mockStepFunctionsClient; + + @BeforeMethod + public void setUp() throws Exception { + MockitoAnnotations.openMocks(this); + stepFunctionService = new StepFunctionService(); + + // Use reflection to set the private field stepFunctionsClient + Field field = StepFunctionService.class.getDeclaredField("stepFunctionsClient"); + field.setAccessible(true); // Make the private field accessible + field.set(stepFunctionService, mockStepFunctionsClient); // Inject mock + } + + @Test + public void testStartExecution_withValidParameters() { + // Arrange + String stateMachineArn = "arn:aws:states:us-east-1:123456789012:stateMachine:exampleStateMachine"; + String inputPayload = "{\"key\":\"value\"}"; + String executionName = "testExecution"; + + StartExecutionResult mockResult = new StartExecutionResult() + .withExecutionArn("arn:aws:states:us-east-1:123456789012:execution:exampleExecution"); + when(mockStepFunctionsClient.startExecution(any(StartExecutionRequest.class))).thenReturn(mockResult); + + // Act + stepFunctionService.startExecution(stateMachineArn, inputPayload, executionName); + + // Assert + ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(StartExecutionRequest.class); + verify(mockStepFunctionsClient).startExecution(requestCaptor.capture()); + + StartExecutionRequest request = requestCaptor.getValue(); + assertEquals(request.getStateMachineArn(), stateMachineArn); + assertEquals(request.getInput(), inputPayload); + assertTrue(request.getName().startsWith("testExecution_")); + } + + @Test(expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = "State Machine ARN cannot be null or empty") + public void testStartExecution_withNullStateMachineArn() { + // Arrange + String stateMachineArn = null; + String inputPayload = "{\"key\":\"value\"}"; + String executionName = "testExecution"; + + // Act + stepFunctionService.startExecution(stateMachineArn, inputPayload, executionName); + } + + @Test(expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = "State Machine ARN cannot be null or empty") + public void testStartExecution_withEmptyStateMachineArn() { + // Arrange + String stateMachineArn = ""; + String inputPayload = "{\"key\":\"value\"}"; + String executionName = "testExecution"; + + // Act + stepFunctionService.startExecution(stateMachineArn, inputPayload, executionName); + } + + @Test + public void testStartExecution_withNullInputPayload() { + // Arrange + String stateMachineArn = "arn:aws:states:us-east-1:123456789012:stateMachine:exampleStateMachine"; + String executionName = "testExecution"; + + StartExecutionResult mockResult = new StartExecutionResult() + .withExecutionArn("arn:aws:states:us-east-1:123456789012:execution:exampleExecution"); + when(mockStepFunctionsClient.startExecution(any(StartExecutionRequest.class))).thenReturn(mockResult); + + // Act + stepFunctionService.startExecution(stateMachineArn, null, executionName); + + // Assert + ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(StartExecutionRequest.class); + verify(mockStepFunctionsClient).startExecution(requestCaptor.capture()); + + StartExecutionRequest request = requestCaptor.getValue(); + assertEquals(request.getStateMachineArn(), stateMachineArn); + assertEquals(request.getInput(), "{}"); // Default to empty payload + } + + @Test + public void testSanitizeExecutionName_withInvalidCharacters() { + // Arrange + String invalidExecutionName = "test/execution:name*with?invalid|characters"; + + // Act + String sanitized = stepFunctionService.sanitizeExecutionName(invalidExecutionName); + + // Assert + assertEquals(sanitized, "test_execution_name_with_invalid_characters"); + } + + @Test + public void testSanitizeExecutionName_withTooLongName() { + // Arrange + String longExecutionName = "ThisIsAVeryLongExecutionNameThatExceedsTheMaximumAllowedLengthOfSixtyNineCharactersAndShouldBeTruncatedAtSomePoint"; + + // Act + String sanitized = stepFunctionService.sanitizeExecutionName(longExecutionName); + + // Assert + assertTrue(sanitized.length() <= 69); + } + + // New Test Cases for Edge Cases + + @Test + public void testSanitizeExecutionName_withValidName() { + // Arrange + String validExecutionName = "validExecutionName"; + + // Act + String sanitized = stepFunctionService.sanitizeExecutionName(validExecutionName); + + // Print the output + System.out.println("Sanitized Execution Name: " + sanitized); + + // Assert + assertEquals(sanitized, validExecutionName); // Should return the same name + } + + @Test + public void testSanitizeExecutionName_withLeadingAndTrailingSpaces() { + // Arrange + String executionName = " executionName "; + + // Act + String sanitized = stepFunctionService.sanitizeExecutionName(executionName); + + // Print the output + System.out.println("Sanitized Execution Name: " + sanitized); + + // Assert + assertEquals(sanitized, "executionName"); // Should trim spaces + } + + @Test(expectedExceptions = IllegalArgumentException.class, + expectedExceptionsMessageRegExp = "Execution name cannot contain only invalid characters") + public void testSanitizeExecutionName_withOnlyInvalidCharacters() { + // Arrange + String invalidOnly = "*/?|<>"; // Input with only invalid characters + + stepFunctionService.sanitizeExecutionName(invalidOnly); + } + + + @Test + public void testSanitizeExecutionName_withMaximumLength() { + // Arrange + String maxLengthExecutionName = "ABCDEFGHIJKLMNOPQRSTUVWXYZABCDEFGHIJKLMNOPQRSTUVWXYZABCDEDHDFHDFHHFCN"; // 69 characters + + // Act + String sanitized = stepFunctionService.sanitizeExecutionName(maxLengthExecutionName); + + // Print the output + System.out.println("Sanitized Execution Name: " + sanitized); + + // Assert + assertEquals(sanitized.length(), 69); // Should be exactly 69 characters + } + + @Test + public void testSanitizeExecutionName_withMultipleInvalidCharacters() { + // Arrange + String executionName = "test//?invalid//name?with*multiple|invalid:characters"; + + // Act + String sanitized = stepFunctionService.sanitizeExecutionName(executionName); + + // Print the output + System.out.println("Sanitized Execution Name: " + sanitized); + + // Assert + assertEquals(sanitized, "test___invalid__name_with_multiple_invalid_characters"); // Should replace all invalid characters + } +} From 1ebfef4d0f73e01bf5a63a78ab51eafc058f416d Mon Sep 17 00:00:00 2001 From: Rishu Yadav Date: Mon, 28 Oct 2024 14:18:52 +0530 Subject: [PATCH 2/2] chore: clear out testing logs --- .../emodb/queue/core/stepfn/StepFunctionServiceTest.java | 8 -------- 1 file changed, 8 deletions(-) diff --git a/queue/src/test/java/com/bazaarvoice/emodb/queue/core/stepfn/StepFunctionServiceTest.java b/queue/src/test/java/com/bazaarvoice/emodb/queue/core/stepfn/StepFunctionServiceTest.java index 09d5fecbe..9ea91d182 100644 --- a/queue/src/test/java/com/bazaarvoice/emodb/queue/core/stepfn/StepFunctionServiceTest.java +++ b/queue/src/test/java/com/bazaarvoice/emodb/queue/core/stepfn/StepFunctionServiceTest.java @@ -135,8 +135,6 @@ public void testSanitizeExecutionName_withValidName() { // Act String sanitized = stepFunctionService.sanitizeExecutionName(validExecutionName); - // Print the output - System.out.println("Sanitized Execution Name: " + sanitized); // Assert assertEquals(sanitized, validExecutionName); // Should return the same name @@ -150,8 +148,6 @@ public void testSanitizeExecutionName_withLeadingAndTrailingSpaces() { // Act String sanitized = stepFunctionService.sanitizeExecutionName(executionName); - // Print the output - System.out.println("Sanitized Execution Name: " + sanitized); // Assert assertEquals(sanitized, "executionName"); // Should trim spaces @@ -175,8 +171,6 @@ public void testSanitizeExecutionName_withMaximumLength() { // Act String sanitized = stepFunctionService.sanitizeExecutionName(maxLengthExecutionName); - // Print the output - System.out.println("Sanitized Execution Name: " + sanitized); // Assert assertEquals(sanitized.length(), 69); // Should be exactly 69 characters @@ -190,8 +184,6 @@ public void testSanitizeExecutionName_withMultipleInvalidCharacters() { // Act String sanitized = stepFunctionService.sanitizeExecutionName(executionName); - // Print the output - System.out.println("Sanitized Execution Name: " + sanitized); // Assert assertEquals(sanitized, "test___invalid__name_with_multiple_invalid_characters"); // Should replace all invalid characters