Search in sources :

Example 1 with Translation

use of org.apache.nifi.processors.yandex.model.Translation in project nifi by apache.

the class TestYandexTranslate method createTestRunner.

private TestRunner createTestRunner(final int statusCode) {
    return TestRunners.newTestRunner(new YandexTranslate() {

        @Override
        protected Invocation prepareResource(final String key, final List<String> text, final String sourceLanguage, final String destLanguage) {
            final Invocation invocation = Mockito.mock(Invocation.class);
            Mockito.doAnswer(new Answer<Response>() {

                @Override
                public Response answer(final InvocationOnMock invocation) throws Throwable {
                    final Response response = Mockito.mock(Response.class);
                    final StatusType statusType = new StatusType() {

                        @Override
                        public int getStatusCode() {
                            return statusCode;
                        }

                        @Override
                        public String getReasonPhrase() {
                            return String.valueOf(statusCode);
                        }

                        @Override
                        public Family getFamily() {
                            return statusCode == 200 ? Family.SUCCESSFUL : Family.SERVER_ERROR;
                        }
                    };
                    Mockito.when(response.getStatus()).thenReturn(statusCode);
                    Mockito.when(response.getStatusInfo()).thenReturn(statusType);
                    if (statusCode == 200) {
                        final Translation translation = new Translation();
                        translation.setCode(statusCode);
                        translation.setLang(destLanguage);
                        final List<String> translationList = new ArrayList<>();
                        for (final String original : text) {
                            final String translated = translations.get(original);
                            translationList.add(translated == null ? original : translated);
                        }
                        translation.setText(translationList);
                        Mockito.when(response.readEntity(Translation.class)).thenReturn(translation);
                    }
                    return response;
                }
            }).when(invocation).invoke();
            return invocation;
        }
    });
}
Also used : Response(javax.ws.rs.core.Response) Answer(org.mockito.stubbing.Answer) Translation(org.apache.nifi.processors.yandex.model.Translation) Invocation(javax.ws.rs.client.Invocation) InvocationOnMock(org.mockito.invocation.InvocationOnMock) StatusType(javax.ws.rs.core.Response.StatusType) ArrayList(java.util.ArrayList)

Example 2 with Translation

use of org.apache.nifi.processors.yandex.model.Translation in project nifi by apache.

the class YandexTranslate method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }
    final StopWatch stopWatch = new StopWatch(true);
    final String key = context.getProperty(KEY).getValue();
    final String sourceLanguage = context.getProperty(SOURCE_LANGUAGE).evaluateAttributeExpressions(flowFile).getValue();
    final String targetLanguage = context.getProperty(TARGET_LANGUAGE).evaluateAttributeExpressions(flowFile).getValue();
    final String encoding = context.getProperty(CHARACTER_SET).evaluateAttributeExpressions(flowFile).getValue();
    final List<String> attributeNames = new ArrayList<>();
    final List<String> textValues = new ArrayList<>();
    for (final PropertyDescriptor descriptor : context.getProperties().keySet()) {
        if (descriptor.isDynamic()) {
            // add to list so that we know the order when the translations come back.
            attributeNames.add(descriptor.getName());
            textValues.add(context.getProperty(descriptor).evaluateAttributeExpressions(flowFile).getValue());
        }
    }
    if (context.getProperty(TRANSLATE_CONTENT).asBoolean()) {
        final byte[] buff = new byte[(int) flowFile.getSize()];
        session.read(flowFile, new InputStreamCallback() {

            @Override
            public void process(final InputStream in) throws IOException {
                StreamUtils.fillBuffer(in, buff);
            }
        });
        final String content = new String(buff, Charset.forName(encoding));
        textValues.add(content);
    }
    final Invocation invocation = prepareResource(key, textValues, sourceLanguage, targetLanguage);
    final Response response;
    try {
        response = invocation.invoke();
    } catch (final Exception e) {
        getLogger().error("Failed to make request to Yandex to transate text for {} due to {}; routing to comms.failure", new Object[] { flowFile, e });
        session.transfer(flowFile, REL_COMMS_FAILURE);
        return;
    }
    if (response.getStatus() != Response.Status.OK.getStatusCode()) {
        getLogger().error("Failed to translate text using Yandex for {}; response was {}: {}; routing to {}", new Object[] { flowFile, response.getStatus(), response.getStatusInfo().getReasonPhrase(), REL_TRANSLATION_FAILED.getName() });
        flowFile = session.putAttribute(flowFile, "yandex.translate.failure.reason", response.getStatusInfo().getReasonPhrase());
        session.transfer(flowFile, REL_TRANSLATION_FAILED);
        return;
    }
    final Map<String, String> newAttributes = new HashMap<>();
    final Translation translation = response.readEntity(Translation.class);
    final List<String> texts = translation.getText();
    for (int i = 0; i < texts.size(); i++) {
        final String text = texts.get(i);
        if (i < attributeNames.size()) {
            final String attributeName = attributeNames.get(i);
            newAttributes.put(attributeName, text);
        } else {
            flowFile = session.write(flowFile, new OutputStreamCallback() {

                @Override
                public void process(final OutputStream out) throws IOException {
                    out.write(text.getBytes(encoding));
                }
            });
            newAttributes.put("language", targetLanguage);
        }
    }
    if (!newAttributes.isEmpty()) {
        flowFile = session.putAllAttributes(flowFile, newAttributes);
    }
    stopWatch.stop();
    session.transfer(flowFile, REL_SUCCESS);
    getLogger().info("Successfully translated {} items for {} from {} to {} in {}; routing to success", new Object[] { texts.size(), flowFile, sourceLanguage, targetLanguage, stopWatch.getDuration() });
}
Also used : FlowFile(org.apache.nifi.flowfile.FlowFile) Translation(org.apache.nifi.processors.yandex.model.Translation) PropertyDescriptor(org.apache.nifi.components.PropertyDescriptor) Invocation(javax.ws.rs.client.Invocation) HashMap(java.util.HashMap) MultivaluedHashMap(javax.ws.rs.core.MultivaluedHashMap) InputStream(java.io.InputStream) OutputStream(java.io.OutputStream) ArrayList(java.util.ArrayList) IOException(java.io.IOException) ProcessException(org.apache.nifi.processor.exception.ProcessException) IOException(java.io.IOException) StopWatch(org.apache.nifi.util.StopWatch) Response(javax.ws.rs.core.Response) InputStreamCallback(org.apache.nifi.processor.io.InputStreamCallback) OutputStreamCallback(org.apache.nifi.processor.io.OutputStreamCallback)

Aggregations

ArrayList (java.util.ArrayList)2 Invocation (javax.ws.rs.client.Invocation)2 Response (javax.ws.rs.core.Response)2 Translation (org.apache.nifi.processors.yandex.model.Translation)2 IOException (java.io.IOException)1 InputStream (java.io.InputStream)1 OutputStream (java.io.OutputStream)1 HashMap (java.util.HashMap)1 MultivaluedHashMap (javax.ws.rs.core.MultivaluedHashMap)1 StatusType (javax.ws.rs.core.Response.StatusType)1 PropertyDescriptor (org.apache.nifi.components.PropertyDescriptor)1 FlowFile (org.apache.nifi.flowfile.FlowFile)1 ProcessException (org.apache.nifi.processor.exception.ProcessException)1 InputStreamCallback (org.apache.nifi.processor.io.InputStreamCallback)1 OutputStreamCallback (org.apache.nifi.processor.io.OutputStreamCallback)1 StopWatch (org.apache.nifi.util.StopWatch)1 InvocationOnMock (org.mockito.invocation.InvocationOnMock)1 Answer (org.mockito.stubbing.Answer)1