/** * Adds a new partition to the working set. */ public void addPartition(PartitionKey partitionKey) { partitions.add(new DefaultConsumablePartition(partitionKey)); }
static DefaultConsumablePartition fromBytes(byte[] bytes) { ByteBuffer bb = ByteBuffer.wrap(bytes); ProcessState processState = ProcessState.fromByte(bb.get()); int keyLength = bb.getInt(); byte[] stringBytes = new byte[keyLength]; bb.get(stringBytes, 0, keyLength); long timestamp = bb.getLong(); int numFailures = bb.getInt(); return new DefaultConsumablePartition(GSON.fromJson(Bytes.toString(stringBytes), PartitionKey.class), processState, timestamp, numFailures); }
@Test(expected = IllegalStateException.class) public void testRetryWithoutTakenTransition() { // cannot take retry a partition without it being taken first ConsumablePartition partition = new DefaultConsumablePartition(generateUniqueKey()); partition.retry(); }
@Test(expected = IllegalStateException.class) public void testAlreadyTakenTransition() { // cannot take a partition that's already taken ConsumablePartition partition = new DefaultConsumablePartition(generateUniqueKey()); partition.take(); partition.take(); }
@Test public void testByteSerialization() { ConsumerWorkingSet workingSet = new ConsumerWorkingSet(); // test with empty partitions lists testSerDe(workingSet); // test with two elements in AVAILABLE and none in IN_PROGRESS workingSet.getPartitions().add(new DefaultConsumablePartition(generateUniqueKey())); workingSet.getPartitions().add(new DefaultConsumablePartition(generateUniqueKey())); testSerDe(workingSet); // test with three elements in partitions and none in inProgressPartitions workingSet.getPartitions().add(new DefaultConsumablePartition(generateUniqueKey())); testSerDe(workingSet); // mark the first element as IN_PROGRESS workingSet.getPartitions().get(0).take(); workingSet.getPartitions().get(0).setTimestamp(System.currentTimeMillis()); testSerDe(workingSet); }
@Test(expected = IllegalStateException.class) public void testAlreadyCompletedTransition() { // cannot complete a partition that has already been completed ConsumablePartition partition = new DefaultConsumablePartition(generateUniqueKey()); partition.take(); partition.complete(); partition.complete(); }
@Test public void testSimpleTransitions() { // tests simple success case ConsumablePartition partition = new DefaultConsumablePartition(generateUniqueKey()); Assert.assertEquals(0, partition.getNumFailures()); partition.take(); partition.retry(); Assert.assertEquals(1, partition.getNumFailures()); partition.take(); // test that untake doesn't increment failure count partition.untake(); Assert.assertEquals(1, partition.getNumFailures()); partition.take(); partition.complete(); }