Examples with WebSocketSession org.apache.nifi.websocket.WebSocketSession used on opensource projects

Example 1 with WebSocketSession

use of org.apache.nifi.websocket.WebSocketSession in project nifi by apache.

the class TestConnectWebSocket method testSuccess.

@Test
public void testSuccess() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(ConnectWebSocket.class);
    final ConnectWebSocket processor = (ConnectWebSocket) runner.getProcessor();
    final SharedSessionState sharedSessionState = new SharedSessionState(processor, new AtomicLong(0));
    // Use this custom session factory implementation so that createdSessions can be read from test case,
    // because MockSessionFactory doesn't expose it.
    final Set<MockProcessSession> createdSessions = new HashSet<>();
    final ProcessSessionFactory sessionFactory = () -> {
        final MockProcessSession session = new MockProcessSession(sharedSessionState, processor);
        createdSessions.add(session);
        return session;
    };
    final WebSocketClientService service = mock(WebSocketClientService.class);
    final WebSocketSession webSocketSession = spy(AbstractWebSocketSession.class);
    when(webSocketSession.getSessionId()).thenReturn("ws-session-id");
    when(webSocketSession.getLocalAddress()).thenReturn(new InetSocketAddress("localhost", 12345));
    when(webSocketSession.getRemoteAddress()).thenReturn(new InetSocketAddress("example.com", 80));
    final String serviceId = "ws-service";
    final String endpointId = "client-1";
    final String textMessageFromServer = "message from server.";
    when(service.getIdentifier()).thenReturn(serviceId);
    when(service.getTargetUri()).thenReturn("ws://example.com/web-socket");
    doAnswer(invocation -> {
        processor.connected(webSocketSession);
        // Two times.
        processor.consume(webSocketSession, textMessageFromServer);
        processor.consume(webSocketSession, textMessageFromServer);
        // Three times.
        final byte[] binaryMessage = textMessageFromServer.getBytes();
        processor.consume(webSocketSession, binaryMessage, 0, binaryMessage.length);
        processor.consume(webSocketSession, binaryMessage, 0, binaryMessage.length);
        processor.consume(webSocketSession, binaryMessage, 0, binaryMessage.length);
        return null;
    }).when(service).connect(endpointId);
    runner.addControllerService(serviceId, service);
    runner.enableControllerService(service);
    runner.setProperty(ConnectWebSocket.PROP_WEBSOCKET_CLIENT_SERVICE, serviceId);
    runner.setProperty(ConnectWebSocket.PROP_WEBSOCKET_CLIENT_ID, endpointId);
    processor.onTrigger(runner.getProcessContext(), sessionFactory);
    final Map<Relationship, List<MockFlowFile>> transferredFlowFiles = getAllTransferredFlowFiles(createdSessions, processor);
    List<MockFlowFile> connectedFlowFiles = transferredFlowFiles.get(AbstractWebSocketGatewayProcessor.REL_CONNECTED);
    assertEquals(1, connectedFlowFiles.size());
    connectedFlowFiles.forEach(ff -> {
        assertFlowFile(webSocketSession, serviceId, endpointId, ff, null);
    });
    List<MockFlowFile> textFlowFiles = transferredFlowFiles.get(AbstractWebSocketGatewayProcessor.REL_MESSAGE_TEXT);
    assertEquals(2, textFlowFiles.size());
    textFlowFiles.forEach(ff -> {
        assertFlowFile(webSocketSession, serviceId, endpointId, ff, WebSocketMessage.Type.TEXT);
    });
    List<MockFlowFile> binaryFlowFiles = transferredFlowFiles.get(AbstractWebSocketGatewayProcessor.REL_MESSAGE_BINARY);
    assertEquals(3, binaryFlowFiles.size());
    binaryFlowFiles.forEach(ff -> {
        assertFlowFile(webSocketSession, serviceId, endpointId, ff, WebSocketMessage.Type.BINARY);
    });
    final List<ProvenanceEventRecord> provenanceEvents = sharedSessionState.getProvenanceEvents();
    assertEquals(6, provenanceEvents.size());
    assertTrue(provenanceEvents.stream().allMatch(event -> ProvenanceEventType.RECEIVE.equals(event.getEventType())));
}

Example 2 with WebSocketSession

use of org.apache.nifi.websocket.WebSocketSession in project nifi by apache.

the class TestPutWebSocket method testServiceIsNotWebSocketService.

@Test
public void testServiceIsNotWebSocketService() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(PutWebSocket.class);
    final ControllerService service = spy(ControllerService.class);
    final WebSocketSession webSocketSession = getWebSocketSession();
    final String serviceId = "ws-service";
    final String endpointId = "client-1";
    final String textMessageFromServer = "message from server.";
    when(service.getIdentifier()).thenReturn(serviceId);
    runner.addControllerService(serviceId, service);
    runner.enableControllerService(service);
    final Map<String, String> attributes = new HashMap<>();
    attributes.put(ATTR_WS_CS_ID, serviceId);
    attributes.put(ATTR_WS_ENDPOINT_ID, endpointId);
    attributes.put(ATTR_WS_SESSION_ID, webSocketSession.getSessionId());
    runner.enqueue(textMessageFromServer, attributes);
    runner.run();
    final List<MockFlowFile> succeededFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_SUCCESS);
    assertEquals(0, succeededFlowFiles.size());
    final List<MockFlowFile> failedFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_FAILURE);
    assertEquals(1, failedFlowFiles.size());
    final MockFlowFile failedFlowFile = failedFlowFiles.iterator().next();
    assertNotNull(failedFlowFile.getAttribute(ATTR_WS_FAILURE_DETAIL));
    final List<ProvenanceEventRecord> provenanceEvents = runner.getProvenanceEvents();
    assertEquals(0, provenanceEvents.size());
}

Example 3 with WebSocketSession

use of org.apache.nifi.websocket.WebSocketSession in project nifi by apache.

the class TestPutWebSocket method testSendFailure.

@Test
public void testSendFailure() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(PutWebSocket.class);
    final WebSocketService service = spy(WebSocketService.class);
    final WebSocketSession webSocketSession = getWebSocketSession();
    final String serviceId = "ws-service";
    final String endpointId = "client-1";
    final String textMessageFromServer = "message from server.";
    when(service.getIdentifier()).thenReturn(serviceId);
    doAnswer(invocation -> {
        final SendMessage sendMessage = invocation.getArgumentAt(2, SendMessage.class);
        sendMessage.send(webSocketSession);
        return null;
    }).when(service).sendMessage(anyString(), anyString(), any(SendMessage.class));
    doThrow(new IOException("Sending message failed.")).when(webSocketSession).sendString(anyString());
    runner.addControllerService(serviceId, service);
    runner.enableControllerService(service);
    final Map<String, String> attributes = new HashMap<>();
    attributes.put(ATTR_WS_CS_ID, serviceId);
    attributes.put(ATTR_WS_ENDPOINT_ID, endpointId);
    attributes.put(ATTR_WS_SESSION_ID, webSocketSession.getSessionId());
    runner.enqueue(textMessageFromServer, attributes);
    runner.run();
    final List<MockFlowFile> succeededFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_SUCCESS);
    assertEquals(0, succeededFlowFiles.size());
    final List<MockFlowFile> failedFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_FAILURE);
    assertEquals(1, failedFlowFiles.size());
    final MockFlowFile failedFlowFile = failedFlowFiles.iterator().next();
    assertNotNull(failedFlowFile.getAttribute(ATTR_WS_FAILURE_DETAIL));
    final List<ProvenanceEventRecord> provenanceEvents = runner.getProvenanceEvents();
    assertEquals(0, provenanceEvents.size());
}

Example 4 with WebSocketSession

use of org.apache.nifi.websocket.WebSocketSession in project nifi by apache.

the class TestPutWebSocket method testSuccess.

@Test
public void testSuccess() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(PutWebSocket.class);
    final WebSocketService service = spy(WebSocketService.class);
    final WebSocketSession webSocketSession = getWebSocketSession();
    final String serviceId = "ws-service";
    final String endpointId = "client-1";
    final String textMessageFromServer = "message from server.";
    when(service.getIdentifier()).thenReturn(serviceId);
    doAnswer(invocation -> {
        final SendMessage sendMessage = invocation.getArgumentAt(2, SendMessage.class);
        sendMessage.send(webSocketSession);
        return null;
    }).when(service).sendMessage(anyString(), anyString(), any(SendMessage.class));
    runner.addControllerService(serviceId, service);
    runner.enableControllerService(service);
    runner.setProperty(PutWebSocket.PROP_WS_MESSAGE_TYPE, "${" + ATTR_WS_MESSAGE_TYPE + "}");
    // Enqueue 1st file as Text.
    final Map<String, String> attributes = new HashMap<>();
    attributes.put(ATTR_WS_CS_ID, serviceId);
    attributes.put(ATTR_WS_ENDPOINT_ID, endpointId);
    attributes.put(ATTR_WS_SESSION_ID, webSocketSession.getSessionId());
    attributes.put(ATTR_WS_MESSAGE_TYPE, WebSocketMessage.Type.TEXT.name());
    runner.enqueue(textMessageFromServer, attributes);
    // Enqueue 2nd file as Binary.
    attributes.put(ATTR_WS_MESSAGE_TYPE, WebSocketMessage.Type.BINARY.name());
    runner.enqueue(textMessageFromServer.getBytes(), attributes);
    runner.run(2);
    final List<MockFlowFile> succeededFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_SUCCESS);
    assertEquals(2, succeededFlowFiles.size());
    assertFlowFile(webSocketSession, serviceId, endpointId, succeededFlowFiles.get(0), WebSocketMessage.Type.TEXT);
    assertFlowFile(webSocketSession, serviceId, endpointId, succeededFlowFiles.get(1), WebSocketMessage.Type.BINARY);
    final List<MockFlowFile> failedFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_FAILURE);
    assertEquals(0, failedFlowFiles.size());
    final List<ProvenanceEventRecord> provenanceEvents = runner.getProvenanceEvents();
    assertEquals(2, provenanceEvents.size());
}

Example 5 with WebSocketSession

use of org.apache.nifi.websocket.WebSocketSession in project nifi by apache.

the class TestPutWebSocket method testServiceIsNotFound.

@Test
public void testServiceIsNotFound() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(PutWebSocket.class);
    final ControllerService service = spy(ControllerService.class);
    final WebSocketSession webSocketSession = getWebSocketSession();
    final String serviceId = "ws-service";
    final String endpointId = "client-1";
    final String textMessageFromServer = "message from server.";
    when(service.getIdentifier()).thenReturn(serviceId);
    runner.addControllerService(serviceId, service);
    runner.enableControllerService(service);
    final Map<String, String> attributes = new HashMap<>();
    attributes.put(ATTR_WS_CS_ID, "different-service-id");
    attributes.put(ATTR_WS_ENDPOINT_ID, endpointId);
    attributes.put(ATTR_WS_SESSION_ID, webSocketSession.getSessionId());
    runner.enqueue(textMessageFromServer, attributes);
    runner.run();
    final List<MockFlowFile> succeededFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_SUCCESS);
    assertEquals(0, succeededFlowFiles.size());
    final List<MockFlowFile> failedFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_FAILURE);
    assertEquals(1, failedFlowFiles.size());
    final MockFlowFile failedFlowFile = failedFlowFiles.iterator().next();
    assertNotNull(failedFlowFile.getAttribute(ATTR_WS_FAILURE_DETAIL));
    final List<ProvenanceEventRecord> provenanceEvents = runner.getProvenanceEvents();
    assertEquals(0, provenanceEvents.size());
}