Parsing multiple web socket onMessage server responses from YAML config

src/integrationTest/java/io/github/azagniotov/stubby4j/yaml/YamlParserTest.java

@@ -1254,7 +1254,7 @@ public void shouldUnmarshall_toWebSocketConfigs() throws Exception {
         final StubWebSocketOnMessageLifeCycle stubWebSocketOnMessageLifeCycle =
                 stubWebSocketConfig.getOnMessage().get(0);
         assertThat(stubWebSocketOnMessageLifeCycle.getClientRequest()).isNotNull();
-        assertThat(stubWebSocketOnMessageLifeCycle.getServerResponse()).isNotNull();
+        assertThat(stubWebSocketOnMessageLifeCycle.getServerResponse(true)).isNotNull();
         assertThat(stubWebSocketOnMessageLifeCycle.getCompleteYAML())
                 .isEqualTo("- client-request:\n" + "    message-type: text\n"
                         + "    body: Hey, server, say apple\n"
@@ -1268,14 +1268,14 @@ public void shouldUnmarshall_toWebSocketConfigs() throws Exception {
         assertThat(clientRequest.getMessageType()).isEqualTo(StubWebSocketMessageType.TEXT);
         assertThat(clientRequest.getBodyAsString()).isEqualTo("Hey, server, say apple");
 
-        final StubWebSocketServerResponse serverResponse = stubWebSocketOnMessageLifeCycle.getServerResponse();
+        final StubWebSocketServerResponse serverResponse = stubWebSocketOnMessageLifeCycle.getServerResponse(true);
         assertThat(serverResponse.getBodyAsString()).isEqualTo("apple");
         assertThat(serverResponse.getDelay()).isEqualTo(500L);
         assertThat(serverResponse.getMessageType()).isEqualTo(StubWebSocketMessageType.TEXT);
         assertThat(serverResponse.getPolicy()).isEqualTo(StubWebSocketServerResponsePolicy.PUSH);
 
         final StubWebSocketServerResponse lastServerResponse =
-                stubWebSocketConfig.getOnMessage().get(2).getServerResponse();
+                stubWebSocketConfig.getOnMessage().get(2).getServerResponse(true);
         final String actualFileContent = "This is response 1 content";
         assertThat(lastServerResponse.getBodyAsString()).isEqualTo(actualFileContent);
         assertThat(lastServerResponse.getBodyAsBytes()).isEqualTo(getBytesUtf8(actualFileContent));
@@ -1346,7 +1346,7 @@ public void shouldUnmarshall_toWebSocketConfigs_withoutOnOpenSection() throws Ex
         final StubWebSocketOnMessageLifeCycle stubWebSocketOnMessageLifeCycle =
                 stubWebSocketConfig.getOnMessage().get(0);
         assertThat(stubWebSocketOnMessageLifeCycle.getClientRequest()).isNotNull();
-        assertThat(stubWebSocketOnMessageLifeCycle.getServerResponse()).isNotNull();
+        assertThat(stubWebSocketOnMessageLifeCycle.getServerResponse(true)).isNotNull();
         assertThat(stubWebSocketOnMessageLifeCycle.getCompleteYAML())
                 .isEqualTo("- client-request:\n" + "    message-type: text\n"
                         + "    body: Hey, server, say apple\n"
@@ -1360,14 +1360,14 @@ public void shouldUnmarshall_toWebSocketConfigs_withoutOnOpenSection() throws Ex
         assertThat(clientRequest.getMessageType()).isEqualTo(StubWebSocketMessageType.TEXT);
         assertThat(clientRequest.getBodyAsString()).isEqualTo("Hey, server, say apple");
 
-        final StubWebSocketServerResponse serverResponse = stubWebSocketOnMessageLifeCycle.getServerResponse();
+        final StubWebSocketServerResponse serverResponse = stubWebSocketOnMessageLifeCycle.getServerResponse(true);
         assertThat(serverResponse.getBodyAsString()).isEqualTo("apple");
         assertThat(serverResponse.getDelay()).isEqualTo(500L);
         assertThat(serverResponse.getMessageType()).isEqualTo(StubWebSocketMessageType.TEXT);
         assertThat(serverResponse.getPolicy()).isEqualTo(StubWebSocketServerResponsePolicy.PUSH);
 
         final StubWebSocketServerResponse lastServerResponse =
-                stubWebSocketConfig.getOnMessage().get(2).getServerResponse();
+                stubWebSocketConfig.getOnMessage().get(2).getServerResponse(true);
         final String actualFileContent = "This is response 1 content";
         assertThat(lastServerResponse.getBodyAsString()).isEqualTo(actualFileContent);
         assertThat(lastServerResponse.getBodyAsBytes()).isEqualTo(getBytesUtf8(actualFileContent));
@@ -1459,6 +1459,146 @@ public void shouldUnmarshall_toWebSocketConfigsWithUuidAndDescription() throws E
         assertThat(stubWebSocketConfig.getOnMessage().isEmpty()).isTrue();
     }
 
+    @Test
+    public void shouldUnmarshall_toWebSocketConfigWithOnMessageSequencedResponses() throws Exception {
+        final URL yamlUrl =
+                YamlParserTest.class.getResource("/yaml/web-socket-valid-config-with-sequenced-responses.yaml");
+        final InputStream stubsConfigStream = yamlUrl.openStream();
+        final String parentDirectory = new File(yamlUrl.getPath()).getParent();
+
+        final YamlParseResultSet yamlParseResultSet =
+                new YamlParser().parse(parentDirectory, inputStreamToString(stubsConfigStream));
+
+        final Map<String, StubWebSocketConfig> webSocketConfigs = yamlParseResultSet.getWebSocketConfigs();
+        assertThat(webSocketConfigs.isEmpty()).isFalse();
+        assertThat(webSocketConfigs.size()).isEqualTo(1);
+
+        final StubWebSocketConfig stubWebSocketConfig =
+                webSocketConfigs.values().iterator().next();
+
+        final StubWebSocketConfig stubWebSocketConfigCopy = webSocketConfigs.get(stubWebSocketConfig.getUrl());
+        assertThat(stubWebSocketConfig).isSameInstanceAs(stubWebSocketConfigCopy);
+        assertThat(stubWebSocketConfig).isEqualTo(stubWebSocketConfigCopy);
+
+        assertThat(stubWebSocketConfig.getDescription()).isEqualTo("this is a web-socket config");
+        assertThat(stubWebSocketConfig.getUrl()).isEqualTo("/items/fruits");
+        assertThat(stubWebSocketConfig.getSubProtocols().size()).isEqualTo(3);
+        assertThat(stubWebSocketConfig.getSubProtocols()).containsExactly("echo", "mamba", "zumba");
+        assertThat(stubWebSocketConfig.getOnOpenServerResponse()).isNotNull();
+        assertThat(stubWebSocketConfig.getOnMessage().size()).isEqualTo(1);
+
+        final StubWebSocketServerResponse onOpenServerResponse = stubWebSocketConfig.getOnOpenServerResponse();
+        assertThat(onOpenServerResponse.getBodyAsString()).isEqualTo("You have been successfully connected");
+        assertThat(onOpenServerResponse.getDelay()).isEqualTo(2000L);
+        assertThat(onOpenServerResponse.getMessageType()).isEqualTo(StubWebSocketMessageType.TEXT);
+        assertThat(onOpenServerResponse.getPolicy()).isEqualTo(StubWebSocketServerResponsePolicy.ONCE);
+
+        final StubWebSocketOnMessageLifeCycle stubWebSocketOnMessageLifeCycle =
+                stubWebSocketConfig.getOnMessage().get(0);
+        assertThat(stubWebSocketOnMessageLifeCycle.getClientRequest()).isNotNull();
+        assertThat(stubWebSocketOnMessageLifeCycle.getServerResponse(false))
+                .isInstanceOf(StubWebSocketServerResponse.class);
+        final String expectedOnMessageConfigAsYAML = ""
+                + "- client-request:\n" + "    message-type: text\n"
+                + "    body: Hey, server, give me fruits\n"
+                + "  server-response:\n"
+                + "  - policy: push\n"
+                + "    message-type: text\n"
+                + "    body: apple\n"
+                + "    delay: 500\n"
+                + "  - policy: push\n"
+                + "    message-type: text\n"
+                + "    body: banana\n"
+                + "    delay: 250\n"
+                + "  - policy: fragmentation\n"
+                + "    message-type: binary\n"
+                + "    file: ../json/response.5.external.file.json\n"
+                + "    delay: 250\n"
+                + "  - policy: push\n"
+                + "    message-type: text\n"
+                + "    body: grapes\n"
+                + "    delay: 500\n";
+        final String actualOnMessageConfigAsYAML = stubWebSocketOnMessageLifeCycle.getCompleteYAML();
+        assertThat(actualOnMessageConfigAsYAML).isEqualTo(expectedOnMessageConfigAsYAML);
+
+        final StubWebSocketClientRequest clientRequest = stubWebSocketOnMessageLifeCycle.getClientRequest();
+        assertThat(clientRequest.getMessageType()).isEqualTo(StubWebSocketMessageType.TEXT);
+        assertThat(clientRequest.getBodyAsString()).isEqualTo("Hey, server, give me fruits");
+
+        final StubWebSocketServerResponse serverResponseSequenceOne =
+                stubWebSocketOnMessageLifeCycle.getServerResponse(true);
+        assertThat(serverResponseSequenceOne.getBodyAsString()).isEqualTo("apple");
+        assertThat(serverResponseSequenceOne.getDelay()).isEqualTo(500L);
+        assertThat(serverResponseSequenceOne.getMessageType()).isEqualTo(StubWebSocketMessageType.TEXT);
+        assertThat(serverResponseSequenceOne.getPolicy()).isEqualTo(StubWebSocketServerResponsePolicy.PUSH);
+
+        final StubWebSocketServerResponse serverResponseSequenceTwo =
+                stubWebSocketOnMessageLifeCycle.getServerResponse(true);
+        assertThat(serverResponseSequenceTwo.getBodyAsString()).isEqualTo("banana");
+        assertThat(serverResponseSequenceTwo.getDelay()).isEqualTo(250L);
+        assertThat(serverResponseSequenceTwo.getMessageType()).isEqualTo(StubWebSocketMessageType.TEXT);
+        assertThat(serverResponseSequenceTwo.getPolicy()).isEqualTo(StubWebSocketServerResponsePolicy.PUSH);
+
+        final StubWebSocketServerResponse serverResponseSequenceThree =
+                stubWebSocketOnMessageLifeCycle.getServerResponse(true);
+        assertThat(serverResponseSequenceThree.getDelay()).isEqualTo(250L);
+        assertThat(serverResponseSequenceThree.getMessageType()).isEqualTo(StubWebSocketMessageType.BINARY);
+        assertThat(serverResponseSequenceThree.getPolicy()).isEqualTo(StubWebSocketServerResponsePolicy.FRAGMENTATION);
+        final String actualFileContent = "" + "mango\n" + "orange\n" + "lemon\n" + "watermelon\n" + "honeydew";
+        assertThat(serverResponseSequenceThree.getBodyAsString()).isEqualTo(actualFileContent);
+        assertThat(serverResponseSequenceThree.getBodyAsBytes()).isEqualTo(getBytesUtf8(actualFileContent));
+
+        final StubWebSocketServerResponse serverResponseSequenceFour =
+                stubWebSocketOnMessageLifeCycle.getServerResponse(true);
+        assertThat(serverResponseSequenceFour.getBodyAsString()).isEqualTo("grapes");
+        assertThat(serverResponseSequenceFour.getDelay()).isEqualTo(500L);
+        assertThat(serverResponseSequenceFour.getMessageType()).isEqualTo(StubWebSocketMessageType.TEXT);
+        assertThat(serverResponseSequenceFour.getPolicy()).isEqualTo(StubWebSocketServerResponsePolicy.PUSH);
+
+        // Loop has restarted - back to first server response, which is "apple"
+        final StubWebSocketServerResponse serverResponseSequenceRestarted =
+                stubWebSocketOnMessageLifeCycle.getServerResponse(true);
+        assertThat(serverResponseSequenceRestarted.getBodyAsString()).isEqualTo("apple");
+        assertThat(serverResponseSequenceRestarted.getDelay()).isEqualTo(500L);
+        assertThat(serverResponseSequenceRestarted.getMessageType()).isEqualTo(StubWebSocketMessageType.TEXT);
+        assertThat(serverResponseSequenceRestarted.getPolicy()).isEqualTo(StubWebSocketServerResponsePolicy.PUSH);
+
+        final String expectedWebSocketConfigAsYAML = ""
+                + "- web-socket:\n"
+                + "    description: this is a web-socket config\n"
+                + "    url: /items/fruits\n"
+                + "    sub-protocols: echo, mamba, zumba\n"
+                + "    on-open:\n"
+                + "      policy: once\n"
+                + "      message-type: text\n"
+                + "      body: You have been successfully connected\n"
+                + "      delay: 2000\n"
+                + "    on-message:\n"
+                + "    - client-request:\n"
+                + "        message-type: text\n"
+                + "        body: Hey, server, give me fruits\n"
+                + "      server-response:\n"
+                + "      - policy: push\n"
+                + "        message-type: text\n"
+                + "        body: apple\n"
+                + "        delay: 500\n"
+                + "      - policy: push\n"
+                + "        message-type: text\n"
+                + "        body: banana\n"
+                + "        delay: 250\n"
+                + "      - policy: fragmentation\n"
+                + "        message-type: binary\n"
+                + "        file: ../json/response.5.external.file.json\n"
+                + "        delay: 250\n"
+                + "      - policy: push\n"
+                + "        message-type: text\n"
+                + "        body: grapes\n"
+                + "        delay: 500\n";
+
+        final String actualWebSocketConfigAsYAML = stubWebSocketConfig.getWebSocketConfigAsYAML();
+        assertThat(actualWebSocketConfigAsYAML).isEqualTo(expectedWebSocketConfigAsYAML);
+    }
+
     @Test
     public void shouldThrowWhenWebSocketConfigWithInvalidServerResponsePolicyName() throws Exception {
 

src/integrationTest/resources/json/response.5.external.file.json

@@ -0,0 +1,5 @@
+mango
+orange
+lemon
+watermelon
+honeydew

src/integrationTest/resources/yaml/web-socket-valid-config-with-sequenced-responses.yaml

@@ -0,0 +1,33 @@
+- web-socket:
+    description: this is a web-socket config
+    url: /items/fruits
+    sub-protocols: echo, mamba, zumba
+
+    on-open:
+      policy: once
+      message-type: text
+      body: You have been successfully connected
+      delay: 2000
+
+    on-message:
+      - client-request:
+          message-type: text
+          body: Hey, server, give me fruits
+
+        server-response:
+          - policy: push
+            message-type: text
+            body: apple
+            delay: 500
+          - policy: push
+            message-type: text
+            body: banana
+            delay: 250
+          - policy: fragmentation
+            message-type: binary
+            file: ../json/response.5.external.file.json
+            delay: 250
+          - policy: push
+            message-type: text
+            body: grapes
+            delay: 500

src/main/java/io/github/azagniotov/stubby4j/server/websocket/StubsServerWebSocket.java

@@ -91,7 +91,7 @@ public void onWebSocketBinary(final byte[] incoming, final int offset, final int
         final List<StubWebSocketOnMessageLifeCycle> onMessage = stubWebSocketConfig.getOnMessage();
         for (final StubWebSocketOnMessageLifeCycle lifeCycle : onMessage) {
             final StubWebSocketClientRequest clientRequest = lifeCycle.getClientRequest();
-            final StubWebSocketServerResponse serverResponse = lifeCycle.getServerResponse();
+            final StubWebSocketServerResponse serverResponse = lifeCycle.getServerResponse(true);
 
             if (Arrays.equals(clientRequest.getBodyAsBytes(), incoming)) {
                 found = true;
@@ -114,7 +114,7 @@ public void onWebSocketText(final String message) {
         final List<StubWebSocketOnMessageLifeCycle> onMessage = stubWebSocketConfig.getOnMessage();
         for (final StubWebSocketOnMessageLifeCycle lifeCycle : onMessage) {
             final StubWebSocketClientRequest clientRequest = lifeCycle.getClientRequest();
-            final StubWebSocketServerResponse serverResponse = lifeCycle.getServerResponse();
+            final StubWebSocketServerResponse serverResponse = lifeCycle.getServerResponse(true);
 
             if (clientRequest.getBodyAsString().equals(message.trim())) {
                 found = true;

src/main/java/io/github/azagniotov/stubby4j/stubs/websocket/StubWebSocketOnMessageLifeCycle.java

@@ -16,20 +16,24 @@
 
 package io.github.azagniotov.stubby4j.stubs.websocket;
 
+import static io.github.azagniotov.generics.TypeSafeConverter.asCheckedLinkedList;
+
 import io.github.azagniotov.stubby4j.annotations.GeneratedCodeMethodCoverageExclusion;
 import io.github.azagniotov.stubby4j.stubs.ReflectableStub;
+import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicInteger;
 
 public class StubWebSocketOnMessageLifeCycle implements ReflectableStub {
 
+    private final AtomicInteger serverResponseSequencedIdCounter = new AtomicInteger(0);
+
     private final StubWebSocketClientRequest clientRequest;
-    private final StubWebSocketServerResponse serverResponse;
+    private final Object serverResponse;
     private final String completeYAML;
 
     public StubWebSocketOnMessageLifeCycle(
-            final StubWebSocketClientRequest clientRequest,
-            final StubWebSocketServerResponse serverResponse,
-            final String completeYAML) {
+            final StubWebSocketClientRequest clientRequest, final Object serverResponse, final String completeYAML) {
         this.clientRequest = clientRequest;
         this.serverResponse = serverResponse;
         this.completeYAML = completeYAML;
@@ -39,8 +43,20 @@ public StubWebSocketClientRequest getClientRequest() {
         return clientRequest;
     }
 
-    public StubWebSocketServerResponse getServerResponse() {
-        return serverResponse;
+    public StubWebSocketServerResponse getServerResponse(final boolean incrementSequencedResponseId) {
+        if (serverResponse instanceof StubWebSocketServerResponse) {
+            return (StubWebSocketServerResponse) serverResponse;
+        }
+
+        final List<StubWebSocketServerResponse> serverResponses =
+                asCheckedLinkedList(this.serverResponse, StubWebSocketServerResponse.class);
+        if (incrementSequencedResponseId) {
+            final int responseSequencedId = serverResponseSequencedIdCounter.getAndIncrement();
+            serverResponseSequencedIdCounter.compareAndSet(serverResponses.size(), 0);
+            return serverResponses.get(responseSequencedId);
+        }
+
+        return serverResponses.get(serverResponseSequencedIdCounter.get());
     }
 
     public String getCompleteYAML() {