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

Example 1 with WebSocketService

use of org.apache.nifi.websocket.WebSocketService 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 2 with WebSocketService

use of org.apache.nifi.websocket.WebSocketService 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 3 with WebSocketService

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

the class PutWebSocket method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession processSession) throws ProcessException {
    final FlowFile flowfile = processSession.get();
    if (flowfile == null) {
        return;
    }
    final String sessionId = context.getProperty(PROP_WS_SESSION_ID).evaluateAttributeExpressions(flowfile).getValue();
    final String webSocketServiceId = context.getProperty(PROP_WS_CONTROLLER_SERVICE_ID).evaluateAttributeExpressions(flowfile).getValue();
    final String webSocketServiceEndpoint = context.getProperty(PROP_WS_CONTROLLER_SERVICE_ENDPOINT).evaluateAttributeExpressions(flowfile).getValue();
    final String messageTypeStr = context.getProperty(PROP_WS_MESSAGE_TYPE).evaluateAttributeExpressions(flowfile).getValue();
    final WebSocketMessage.Type messageType = WebSocketMessage.Type.valueOf(messageTypeStr);
    if (StringUtils.isEmpty(sessionId)) {
        getLogger().debug("Specific SessionID not specified. Message will be broadcast to all connected clients.");
    }
    if (StringUtils.isEmpty(webSocketServiceId) || StringUtils.isEmpty(webSocketServiceEndpoint)) {
        transferToFailure(processSession, flowfile, "Required WebSocket attribute was not found.");
        return;
    }
    final ControllerService controllerService = context.getControllerServiceLookup().getControllerService(webSocketServiceId);
    if (controllerService == null) {
        transferToFailure(processSession, flowfile, "WebSocket ControllerService was not found.");
        return;
    } else if (!(controllerService instanceof WebSocketService)) {
        transferToFailure(processSession, flowfile, "The ControllerService found was not a WebSocket ControllerService but a " + controllerService.getClass().getName());
        return;
    }
    final WebSocketService webSocketService = (WebSocketService) controllerService;
    final byte[] messageContent = new byte[(int) flowfile.getSize()];
    final long startSending = System.currentTimeMillis();
    final AtomicReference<String> transitUri = new AtomicReference<>();
    final Map<String, String> attrs = new HashMap<>();
    attrs.put(ATTR_WS_CS_ID, webSocketService.getIdentifier());
    if (!StringUtils.isEmpty(sessionId)) {
        attrs.put(ATTR_WS_SESSION_ID, sessionId);
    }
    attrs.put(ATTR_WS_ENDPOINT_ID, webSocketServiceEndpoint);
    attrs.put(ATTR_WS_MESSAGE_TYPE, messageTypeStr);
    processSession.read(flowfile, in -> {
        StreamUtils.fillBuffer(in, messageContent, true);
    });
    try {
        webSocketService.sendMessage(webSocketServiceEndpoint, sessionId, sender -> {
            switch(messageType) {
                case TEXT:
                    sender.sendString(new String(messageContent, CHARSET_NAME));
                    break;
                case BINARY:
                    sender.sendBinary(ByteBuffer.wrap(messageContent));
                    break;
            }
            attrs.put(ATTR_WS_LOCAL_ADDRESS, sender.getLocalAddress().toString());
            attrs.put(ATTR_WS_REMOTE_ADDRESS, sender.getRemoteAddress().toString());
            transitUri.set(sender.getTransitUri());
        });
        final FlowFile updatedFlowFile = processSession.putAllAttributes(flowfile, attrs);
        final long transmissionMillis = System.currentTimeMillis() - startSending;
        processSession.getProvenanceReporter().send(updatedFlowFile, transitUri.get(), transmissionMillis);
        processSession.transfer(updatedFlowFile, REL_SUCCESS);
    } catch (WebSocketConfigurationException | IllegalStateException | IOException e) {
        // WebSocketConfigurationException: If the corresponding WebSocketGatewayProcessor has been stopped.
        // IllegalStateException: Session is already closed or not found.
        // IOException: other IO error.
        getLogger().error("Failed to send message via WebSocket due to " + e, e);
        transferToFailure(processSession, flowfile, e.toString());
    }
}

Example 4 with WebSocketService

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

the class TestPutWebSocket method testSessionIsNotSpecified.

@Test
public void testSessionIsNotSpecified() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(PutWebSocket.class);
    final WebSocketService service = spy(WebSocketService.class);
    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);
    runner.enqueue(textMessageFromServer, attributes);
    runner.run();
    final List<MockFlowFile> succeededFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_SUCCESS);
    // assertEquals(0, succeededFlowFiles.size());   //No longer valid test after NIFI-3318 since not specifying sessionid will send to all clients
    assertEquals(1, succeededFlowFiles.size());
    final List<MockFlowFile> failedFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_FAILURE);
    // assertEquals(1, failedFlowFiles.size());      //No longer valid test after NIFI-3318
    assertEquals(0, failedFlowFiles.size());
    final List<ProvenanceEventRecord> provenanceEvents = runner.getProvenanceEvents();
    assertEquals(0, provenanceEvents.size());
}