@Override
public MockFlowFile penalize(final FlowFile flowFile) {
validateState(flowFile);
final MockFlowFile mockFlowFile = (MockFlowFile) flowFile;
mockFlowFile.setPenalized();
return mockFlowFile;
}
@Override
public MockFlowFile importFrom(
final Path path, final boolean keepSourceFile, final FlowFile flowFile) {
validateState(flowFile);
if (path == null || flowFile == null) {
throw new IllegalArgumentException("argument cannot be null");
}
if (!(flowFile instanceof MockFlowFile)) {
throw new IllegalArgumentException("Cannot export a flow file that I did not create");
}
final MockFlowFile mock = (MockFlowFile) flowFile;
MockFlowFile newFlowFile = new MockFlowFile(mock.getId(), flowFile);
currentVersions.put(newFlowFile.getId(), newFlowFile);
final ByteArrayOutputStream baos = new ByteArrayOutputStream();
try {
Files.copy(path, baos);
} catch (final IOException e) {
throw new FlowFileAccessException(e.toString(), e);
}
newFlowFile.setData(baos.toByteArray());
newFlowFile =
putAttribute(newFlowFile, CoreAttributes.FILENAME.key(), path.getFileName().toString());
return newFlowFile;
}
@Override
public MockFlowFile write(final FlowFile flowFile, final StreamCallback callback) {
validateState(flowFile);
if (callback == null || flowFile == null) {
throw new IllegalArgumentException("argument cannot be null");
}
if (!(flowFile instanceof MockFlowFile)) {
throw new IllegalArgumentException("Cannot export a flow file that I did not create");
}
final MockFlowFile mock = (MockFlowFile) flowFile;
final ByteArrayInputStream in = new ByteArrayInputStream(mock.getData());
final ByteArrayOutputStream out = new ByteArrayOutputStream();
try {
callback.process(in, out);
} catch (final IOException e) {
throw new ProcessException(e.toString(), e);
}
final MockFlowFile newFlowFile = new MockFlowFile(mock.getId(), flowFile);
currentVersions.put(newFlowFile.getId(), newFlowFile);
newFlowFile.setData(out.toByteArray());
return newFlowFile;
}
@Override
public MockFlowFile putAttribute(
final FlowFile flowFile, final String attrName, final String attrValue) {
validateState(flowFile);
if (attrName == null || attrValue == null || flowFile == null) {
throw new IllegalArgumentException("argument cannot be null");
}
if (!(flowFile instanceof MockFlowFile)) {
throw new IllegalArgumentException(
"Cannot update attributes of a flow file that I did not create");
}
if ("uuid".equals(attrName)) {
Assert.fail(
"Should not be attempting to set FlowFile UUID via putAttribute. This will be ignored in production");
}
final MockFlowFile mock = (MockFlowFile) flowFile;
final MockFlowFile newFlowFile = new MockFlowFile(mock.getId(), flowFile);
currentVersions.put(newFlowFile.getId(), newFlowFile);
final Map<String, String> attrs = new HashMap<>();
attrs.put(attrName, attrValue);
newFlowFile.putAttributes(attrs);
return newFlowFile;
}
@Override
public MockFlowFile clone(final FlowFile flowFile) {
validateState(flowFile);
final MockFlowFile newFlowFile = new MockFlowFile(sharedState.nextFlowFileId(), flowFile);
currentVersions.put(newFlowFile.getId(), newFlowFile);
beingProcessed.add(newFlowFile.getId());
return newFlowFile;
}
@Override
public void remove(final Collection<FlowFile> flowFiles) {
for (final FlowFile flowFile : flowFiles) {
validateState(flowFile);
}
for (final FlowFile flowFile : flowFiles) {
remove(flowFile);
}
}
@Override
public void transfer(final FlowFile flowFile) {
validateState(flowFile);
if (!(flowFile instanceof MockFlowFile)) {
throw new IllegalArgumentException("I only accept MockFlowFile");
}
beingProcessed.remove(flowFile.getId());
processorQueue.offer((MockFlowFile) flowFile);
}
@Override
public void transfer(final FlowFile flowFile, final Relationship relationship) {
validateState(flowFile);
List<MockFlowFile> list = transferMap.get(relationship);
if (list == null) {
list = new ArrayList<>();
transferMap.put(relationship, list);
}
beingProcessed.remove(flowFile.getId());
list.add((MockFlowFile) flowFile);
}
@Override
public MockFlowFile merge(
Collection<FlowFile> sources,
FlowFile destination,
byte[] header,
byte[] footer,
byte[] demarcator) {
for (final FlowFile flowFile : sources) {
validateState(flowFile);
}
validateState(destination);
final ByteArrayOutputStream baos = new ByteArrayOutputStream();
try {
if (header != null) {
baos.write(header);
}
int count = 0;
for (final FlowFile flowFile : sources) {
baos.write(((MockFlowFile) flowFile).getData());
if (demarcator != null && ++count != sources.size()) {
baos.write(demarcator);
}
}
if (footer != null) {
baos.write(footer);
}
} catch (final IOException e) {
throw new AssertionError("failed to write data to BAOS");
}
final MockFlowFile newFlowFile = new MockFlowFile(destination.getId(), destination);
newFlowFile.setData(baos.toByteArray());
currentVersions.put(newFlowFile.getId(), newFlowFile);
return newFlowFile;
}
@Override
public MockFlowFile merge(final Collection<FlowFile> sources, final FlowFile destination) {
for (final FlowFile source : sources) {
validateState(source);
}
validateState(destination);
final ByteArrayOutputStream baos = new ByteArrayOutputStream();
for (final FlowFile flowFile : sources) {
final MockFlowFile mock = (MockFlowFile) flowFile;
final byte[] data = mock.getData();
try {
baos.write(data);
} catch (final IOException e) {
throw new AssertionError("Failed to write to BAOS");
}
}
final MockFlowFile newFlowFile = new MockFlowFile(destination.getId(), destination);
newFlowFile.setData(baos.toByteArray());
currentVersions.put(newFlowFile.getId(), newFlowFile);
return newFlowFile;
}
@Override
public MockFlowFile removeAllAttributes(final FlowFile flowFile, final Set<String> attrNames) {
validateState(flowFile);
if (attrNames == null || flowFile == null) {
throw new IllegalArgumentException("argument cannot be null");
}
if (!(flowFile instanceof MockFlowFile)) {
throw new IllegalArgumentException("Cannot export a flow file that I did not create");
}
final MockFlowFile mock = (MockFlowFile) flowFile;
final MockFlowFile newFlowFile = new MockFlowFile(mock.getId(), flowFile);
currentVersions.put(newFlowFile.getId(), newFlowFile);
newFlowFile.removeAttributes(attrNames);
return newFlowFile;
}
@Override
public void remove(final FlowFile flowFile) {
validateState(flowFile);
final Iterator<Long> itr = beingProcessed.iterator();
while (itr.hasNext()) {
final Long ffId = itr.next();
if (ffId != null && ffId.equals(flowFile.getId())) {
itr.remove();
beingProcessed.remove(ffId);
removedCount++;
currentVersions.remove(ffId);
return;
}
}
throw new ProcessException(flowFile + " not found in queue");
}
@Override
public MockFlowFile removeAllAttributes(final FlowFile flowFile, final Pattern keyPattern) {
validateState(flowFile);
if (flowFile == null) {
throw new IllegalArgumentException("flowFile cannot be null");
}
if (keyPattern == null) {
return (MockFlowFile) flowFile;
}
final Set<String> attrsToRemove = new HashSet<>();
for (final String key : flowFile.getAttributes().keySet()) {
if (keyPattern.matcher(key).matches()) {
attrsToRemove.add(key);
}
}
return removeAllAttributes(flowFile, attrsToRemove);
}
@Override
public void read(final FlowFile flowFile, final InputStreamCallback callback) {
if (callback == null || flowFile == null) {
throw new IllegalArgumentException("argument cannot be null");
}
validateState(flowFile);
if (!(flowFile instanceof MockFlowFile)) {
throw new IllegalArgumentException("Cannot export a flow file that I did not create");
}
final MockFlowFile mock = (MockFlowFile) flowFile;
final ByteArrayInputStream bais = new ByteArrayInputStream(mock.getData());
try {
callback.process(bais);
} catch (final IOException e) {
throw new ProcessException(e.toString(), e);
}
}
@Override
public void exportTo(final FlowFile flowFile, final OutputStream out) {
validateState(flowFile);
if (flowFile == null || out == null) {
throw new IllegalArgumentException("arguments cannot be null");
}
if (!(flowFile instanceof MockFlowFile)) {
throw new IllegalArgumentException("Cannot export a flow file that I did not create");
}
final MockFlowFile mock = (MockFlowFile) flowFile;
try {
out.write(mock.getData());
} catch (final IOException e) {
throw new FlowFileAccessException(e.toString(), e);
}
}
@Override
public void exportTo(final FlowFile flowFile, final Path path, final boolean append) {
validateState(flowFile);
if (flowFile == null || path == null) {
throw new IllegalArgumentException("argument cannot be null");
}
if (!(flowFile instanceof MockFlowFile)) {
throw new IllegalArgumentException("Cannot export a flow file that I did not create");
}
final MockFlowFile mock = (MockFlowFile) flowFile;
final OpenOption mode = append ? StandardOpenOption.APPEND : StandardOpenOption.CREATE;
try (final OutputStream out = Files.newOutputStream(path, mode)) {
out.write(mock.getData());
} catch (final IOException e) {
throw new FlowFileAccessException(e.toString(), e);
}
}
@Override
public MockFlowFile importFrom(final InputStream in, final FlowFile flowFile) {
validateState(flowFile);
if (in == null || flowFile == null) {
throw new IllegalArgumentException("argument cannot be null");
}
if (!(flowFile instanceof MockFlowFile)) {
throw new IllegalArgumentException("Cannot export a flow file that I did not create");
}
final MockFlowFile mock = (MockFlowFile) flowFile;
final MockFlowFile newFlowFile = new MockFlowFile(mock.getId(), flowFile);
currentVersions.put(newFlowFile.getId(), newFlowFile);
try {
final byte[] data = readFully(in);
newFlowFile.setData(data);
return newFlowFile;
} catch (final IOException e) {
throw new FlowFileAccessException(e.toString(), e);
}
}
@Override
public MockFlowFile clone(final FlowFile flowFile, final long offset, final long size) {
validateState(flowFile);
if (offset + size > flowFile.getSize()) {
throw new FlowFileHandlingException(
"Specified offset of "
+ offset
+ " and size "
+ size
+ " exceeds size of "
+ flowFile.toString());
}
final MockFlowFile newFlowFile = new MockFlowFile(sharedState.nextFlowFileId(), flowFile);
final byte[] newContent =
Arrays.copyOfRange(
((MockFlowFile) flowFile).getData(), (int) offset, (int) (offset + size));
newFlowFile.setData(newContent);
currentVersions.put(newFlowFile.getId(), newFlowFile);
beingProcessed.add(newFlowFile.getId());
return newFlowFile;
}
public byte[] getContentAsByteArray(final MockFlowFile flowFile) {
validateState(flowFile);
return flowFile.getData();
}
