Repository: apache/kafka
File: connect/runtime/src/test/java/org/apache/kafka/connect/storage/KafkaStatusBackingStoreTest.java (Java)

Code before:

    @Test
    public void deleteConnectorState() {
        final byte[] value = new byte[0];
        Map<String, Object> statusMap = new HashMap<>();
        statusMap.put("worker_id", WORKER_ID);
        statusMap.put("state", "RUNNING");
        statusMap.put("generation", 0L);

        when(converter.fromConnectData(eq(STATUS_TOPIC), any(Schema.class), any(Struct.class))).thenReturn(value);
        when(converter.fromConnectData(eq(STATUS_TOPIC), any(Schema.class), any(Struct.class))).thenReturn(value);
        when(converter.toConnectData(STATUS_TOPIC, value)).thenReturn(new SchemaAndValue(null, statusMap));

        ConnectorStatus connectorStatus = new ConnectorStatus(CONNECTOR, ConnectorStatus.State.RUNNING, WORKER_ID, 0);
        store.put(connectorStatus);
        TaskStatus taskStatus = new TaskStatus(TASK, TaskStatus.State.RUNNING, WORKER_ID, 0);
        store.put(taskStatus);
        store.read(consumerRecord(0, "status-task-conn-0", value));

        verify(kafkaBasedLog).send(eq("status-connector-" + CONNECTOR), eq(value), any(Callback.class));
        verify(kafkaBasedLog).send(eq("status-task-conn-0"), eq(value), any(Callback.class));

        assertEquals(new HashSet<>(Collections.singletonList(CONNECTOR)), store.connectors());
        assertEquals(new HashSet<>(Collections.singletonList(taskStatus)), new HashSet<>(store.getAll(CONNECTOR)));
        store.read(consumerRecord(0, "status-connector-conn", null));
        assertTrue(store.connectors().isEmpty());
        assertTrue(store.getAll(CONNECTOR).isEmpty());
    }

    @Test
    public void deleteTaskState() {
        final byte[] value = new byte[0];
        Map<String, Object> statusMap = new HashMap<>();
        statusMap.put("worker_id", WORKER_ID);
        statusMap.put("state", "RUNNING");
        statusMap.put("generation", 0L);

        when(converter.fromConnectData(eq(STATUS_TOPIC), any(Schema.class), any(Struct.class))).thenReturn(value);
        when(converter.toConnectData(STATUS_TOPIC, value)).thenReturn(new SchemaAndValue(null, statusMap));

        TaskStatus taskStatus = new TaskStatus(TASK, TaskStatus.State.RUNNING, WORKER_ID, 0);
        store.put(taskStatus);
        store.read(consumerRecord(0, "status-task-conn-0", value));

        verify(kafkaBasedLog).send(eq("status-task-conn-0"), eq(value), any(Callback.class));

        assertEquals(new HashSet<>(Collections.singletonList(taskStatus)), new HashSet<>(store.getAll(CONNECTOR)));
        store.read(consumerRecord(0, "status-task-conn-0", null));
        assertTrue(store.getAll(CONNECTOR).isEmpty());
    }


Review comment (suggestion):
I'm not sure what the point of this verification is since we're only making calls to the store's read methods and this verification would always be true with or without your other changes.

Code after:

    @Test
    public void readTaskStateShouldIgnoreStaleStatusesFromOtherWorkers() {
        byte[] value = new byte[0];
        String otherWorkerId = "anotherhost:8083";
        String yetAnotherWorkerId = "yetanotherhost:8083";

        // This worker sends a RUNNING status in the most recent generation
        Map<String, Object> firstStatusRead = new HashMap<>();
        firstStatusRead.put("worker_id", otherWorkerId);
        firstStatusRead.put("state", "RUNNING");
        firstStatusRead.put("generation", 10L);

        // Another worker still ends up producing an UNASSIGNED status before it could
        // read the newer RUNNING status from above belonging to an older generation.
        Map<String, Object> secondStatusRead = new HashMap<>();
        secondStatusRead.put("worker_id", WORKER_ID);
        secondStatusRead.put("state", "UNASSIGNED");
        secondStatusRead.put("generation", 9L);

        Map<String, Object> thirdStatusRead = new HashMap<>();
        thirdStatusRead.put("worker_id", yetAnotherWorkerId);
        thirdStatusRead.put("state", "RUNNING");
        thirdStatusRead.put("generation", 1L);

        when(converter.toConnectData(STATUS_TOPIC, value))
                .thenReturn(new SchemaAndValue(null, firstStatusRead))
                .thenReturn(new SchemaAndValue(null, secondStatusRead))
                .thenReturn(new SchemaAndValue(null, thirdStatusRead));

        when(converter.fromConnectData(eq(STATUS_TOPIC), any(Schema.class), any(Struct.class)))
                .thenReturn(value);

        doAnswer(invocation -> {
            ((Callback) invocation.getArgument(2)).onCompletion(null, null);
            store.read(consumerRecord(2, "status-task-conn-0", value));
            return null;
        }).when(kafkaBasedLog).send(eq("status-task-conn-0"), eq(value), any(Callback.class));

        store.read(consumerRecord(0, "status-task-conn-0", value));
        store.read(consumerRecord(1, "status-task-conn-0", value));

        // The latest task status should reflect RUNNING status from the newer generation
        TaskStatus status = new TaskStatus(TASK, TaskStatus.State.RUNNING, otherWorkerId, 10);
        assertEquals(status, store.get(TASK));

        // This  status is from the another worker not necessarily belonging to the above group.
        // In this case, the status should get updated irrespective of whatever status info was present before.
        TaskStatus latestStatus = new TaskStatus(TASK, TaskStatus.State.RUNNING, yetAnotherWorkerId, 1);
        store.put(latestStatus);
        assertEquals(latestStatus, store.get(TASK));