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;
}
});
}
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() });
}
Aggregations