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