@Override
public List<FlowFile> get(final FlowFileFilter filter) {
final List<FlowFile> flowFiles = new ArrayList<>();
final List<MockFlowFile> unselected = new ArrayList<>();
while (true) {
final MockFlowFile flowFile = processorQueue.poll();
if (flowFile == null) {
break;
}
final FlowFileFilter.FlowFileFilterResult filterResult = filter.filter(flowFile);
if (filterResult.isAccept()) {
flowFiles.add(flowFile);
beingProcessed.add(flowFile.getId());
currentVersions.put(flowFile.getId(), flowFile);
originalVersions.put(flowFile.getId(), flowFile);
} else {
unselected.add(flowFile);
}
if (!filterResult.isContinue()) {
break;
}
}
processorQueue.addAll(unselected);
return flowFiles;
}
@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 MockFlowFile get() {
final MockFlowFile flowFile = processorQueue.poll();
if (flowFile != null) {
beingProcessed.add(flowFile.getId());
currentVersions.put(flowFile.getId(), flowFile);
originalVersions.put(flowFile.getId(), flowFile);
}
return flowFile;
}
@Override
public void rollback(final boolean penalize) {
for (final List<MockFlowFile> list : transferMap.values()) {
for (final MockFlowFile flowFile : list) {
processorQueue.offer(flowFile);
}
}
for (final Long flowFileId : beingProcessed) {
final MockFlowFile flowFile = originalVersions.get(flowFileId);
if (flowFile != null) {
processorQueue.offer(flowFile);
}
}
rolledback = true;
beingProcessed.clear();
currentVersions.clear();
originalVersions.clear();
transferMap.clear();
clearTransferState();
}
@Override
public QueueSize getQueueSize() {
return flowFileQueue.size();
}
@Override
public void enqueue(final FlowFile... flowFiles) {
for (final FlowFile flowFile : flowFiles) {
flowFileQueue.offer((MockFlowFile) flowFile);
}
}
@Override
public void assertQueueNotEmpty() {
Assert.assertFalse(flowFileQueue.isEmpty());
}
@Override
public void assertQueueEmpty() {
Assert.assertTrue(flowFileQueue.isEmpty());
}
@Override
public boolean isQueueEmpty() {
return flowFileQueue.isEmpty();
}
@Override
public QueueSize getQueueSize() {
return processorQueue.size();
}