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