@Override
      public Void answer(InvocationOnMock invocation) throws Throwable {
        SystemStreamPartition ssp = invocation.getArgumentAt(0, SystemStreamPartition.class);
        if (partitionOfFirstMessage == -1) {
          partitionOfFirstMessage = ssp.getPartition().getPartitionId();
        } else {
          partitionOfSecondMessage = ssp.getPartition().getPartitionId();
          Assert.assertEquals(partitionOfFirstMessage, partitionOfSecondMessage);
          testFinished = true;
        }
        return null;
      }
Test keys of type byte[] goes to the same partition if they have the same contents.
samza-core/src/test/java/org/apache/samza/system/inmemory/TestInMemorySystemProducer.java