Тестирование с TestComputationHarness в YTsaurus Flow (Java)

TestComputationHarness harness = TestComputationHarness.builder()
        .setPipelineContext(pipelineContext)       // обязательно
        .setPipelineSpec(txtSpec)                  // обязательно: String, YTreeNode или InputStream
        .addExternalStateSchema("state-name", schema) // опционально
        .setJobContext(jobContext)                 // опционально, по умолчанию timeout=10 мин
        .build();

var request = TestDoProcessRequest.builder("computation-id")
        .setMessages(messages)           // List<ExtendedMessage>
        .setTimers(timers)               // List<Timer>
        .setWatermarks(watermarks)       // Map<String, Long>
        .setExternalState("name", stateMap)  // Map<Payload, ExternalState>
        .setInternalState("name", stateMap)  // Map<Payload, InternalState>
        .build();

ExtendedMessage msg = ExtendedMessage.builder()
        .setStreamId("hit")
        .setKey(key)
        .setPayload(new Hit("hit-1", 1000L, "payload-data"))
        .setEventTimestamp(1000L)
        .build();

var hitSchema = testHarness.getStream("hit").getSchema();
var payload = new PayloadBuilder(hitSchema)
        .set("hit_id", "hit-1")
        .set("hit_time", 1000L)
        .set("hit_payload", "payload-data")
        .finish();

ExtendedMessage msg = ExtendedMessage.builder()
        .setStreamId("hit")
        .setKey(key)
        .setPayload(payload)
        .setEventTimestamp(1000L)
        .build();

TableSchema keySchema = testHarness.getGroupBySchema("join");

Payload key = new PayloadBuilder(keySchema)
        .set("hash", 0L)
        .set("hit_id", "hit-1")
        .set("hit_time", 1000L)
        .finish();

Timer timer = Timer.builder()
        .setStreamId("timer")
        .setKey(key)
        .setTriggerTimestamp(triggerTimestamp)
        .build();

Map<String, Long> watermarks = Map.of(
        "hit", 0L,
        "action", 0L
);

TableSchema stateSchema = TableSchema.builder()
        .addValue("hit_payload", ColumnValueType.STRING)
        .addValue("show_time", ColumnValueType.UINT64)
        .addValue("click_time", ColumnValueType.UINT64)
        .build();

ExternalState preState = new ExternalState(
        new PayloadBuilder(stateSchema)
                .set("hit_payload", "some-payload")
                .set("show_time", showTime)
                .set("click_time", clickTime)
                .finish()
);

var request = TestDoProcessRequest.builder("join")
        .setTimers(List.of(timer))
        .setExternalState("join-state", Map.of(key, preState))
        .build();

TJoinState protoState = TJoinState.newBuilder()
        .setHitPayload("some-payload")
        .setShowTime(showTime)
        .setClickTime(clickTime)
        .build();

InternalState internalState = new InternalState(protoState.toByteArray());

var request = TestDoProcessRequest.builder("join")
        .setTimers(List.of(timer))
        .setInternalState("join-state", Map.of(key, internalState))
        .build();

var response = testHarness.doProcess(request);

// Проверяем количество выходных сообщений
assertEquals(1, response.getOutputMessagesFlatten().size());

// Получаем сообщение и проверяем стрим
var msg = response.getOutputMessagesFlatten().get(0);
assertEquals("joined_action", msg.getStreamId());

// Для типизированных стримов — приведение payload к POJO
JoinedAction result = (JoinedAction) msg.getPayload();
assertEquals("hit-1", result.getHitId());

// Для нетипизированных стримов — чтение полей через get()
byte[] data = msg.get("data", byte[].class);

assertEquals(1, response.getOutputTimersFlatten().size());
var timer = response.getOutputTimersFlatten().get(0);
assertEquals(expectedTriggerTimestamp, timer.getTriggerTimestamp());

// External state
assertEquals(1, response.getExternalStateSize("join-state"));
var state = response.getExternalState("join-state", key);
assertFalse(state.isReset());
assertEquals("payload-data", state.getValue().get("hit_payload", String.class));

// Проверка сброса состояния
var stateAfterTimer = response.getExternalState("join-state", key);
assertTrue(stateAfterTimer.isReset());

// Internal state (Proto)
var internalState = response.getInternalState("join-state", key);
assertFalse(internalState.isReset());
TJoinState joinState = ProtoUtils.parseBytes(internalState.getValue(), TJoinState.class);
assertEquals("payload-data", joinState.getHitPayload());

// Полная картина: загруженное состояние видно, даже если не изменялось
assertEquals(1, response.getExternalStateSize("join-state"));
assertNotNull(response.getExternalState("join-state", key));

// Только изменённые состояния: computation ничего не записал для этого ключа
assertEquals(0, response.getModifiedExternalStateSize("join-state"));
assertNull(response.getModifiedExternalState("join-state", key));

// Неизвестные имя/ключ не бросают исключение
assertNull(response.getExternalState("unknown", key));
assertTrue(response.getExternalStates("unknown").isEmpty());

public class JoinProcessFunctionTest {

    private static final long WAIT_SECONDS = 10L;
    private static final long BASE_HIT_TIME = 1000L;
    private static final long WATERMARK = 0L;

    private TestComputationHarness testHarness;
    private TableSchema keySchema;
    private TableSchema joinStateSchema;

    @BeforeEach
    void init() throws IOException {
        // 1. Создаём PipelineContext вручную
        var pipelineContext = new PipelineContext();

        // 2. Регистрируем computation с process function
        Computation join = Computation.builder()
                .setComputationId("join")
                .setProcessFunction(new JoinProcessFunction())
                .build();
        pipelineContext.registerComputation(join);

        // 3. Регистрируем типизированные стримы
        pipelineContext.registerStream(FlowStreams.typed("hit", Hit.class));
        pipelineContext.registerStream(FlowStreams.typed("action", Action.class));
        pipelineContext.registerStream(FlowStreams.typed("joined_action", JoinedAction.class));

        // 4. Определяем схему внешнего состояния
        this.joinStateSchema = TableSchema.builder()
                .addValue("hit_payload", ColumnValueType.STRING)
                .addValue("show_time", ColumnValueType.UINT64)
                .addValue("click_time", ColumnValueType.UINT64)
                .build();

        // 5. Читаем pipeline spec и создаём harness
        var specPath = Paths.getSourcePath(
                "yt/yt/flow/examples/java/wait_click_join/test/pipeline.yson");
        var txtSpec = Files.readString(Path.of(specPath));

        this.testHarness = TestComputationHarness.builder()
                .setPipelineContext(pipelineContext)
                .setPipelineSpec(txtSpec)
                .addExternalStateSchema("join-state", joinStateSchema)
                .build();

        // 6. Получаем key schema для join computation (group_by_schema из pipeline.yson)
        this.keySchema = testHarness.getGroupBySchema("join");
    }

    // — Вспомогательные методы —

    private Payload buildKey(String hitId, long hitTime) {
        return new PayloadBuilder(keySchema)
                .set("hash", 0L)
                .set("hit_id", hitId)
                .set("hit_time", hitTime)
                .finish();
    }

    private ExtendedMessage buildHitMessage(String hitId, long hitTime, String hitPayload) {
        return ExtendedMessage.builder()
                .setStreamId("hit")
                .setKey(buildKey(hitId, hitTime))
                .setPayload(new Hit(hitId, hitTime, hitPayload))
                .setEventTimestamp(hitTime)
                .build();
    }

    private Map<String, Long> defaultWatermarks() {
        return Map.of("hit", WATERMARK, "action", WATERMARK);
    }

    // — Тесты —

    @Test
    void testHitMessageStoresHitPayloadInState() {
        var messages = List.of(buildHitMessage("hit-1", BASE_HIT_TIME, "payload-data"));

        var request = TestDoProcessRequest.builder("join")
                .setMessages(messages)
                .setWatermarks(defaultWatermarks())
                .build();

        var response = testHarness.doProcess(request);

        // Нет выходных сообщений — результат будет при срабатывании таймера
        assertTrue(response.getOutputMessagesFlatten().isEmpty());

        // Таймер должен быть установлен
        assertEquals(1, response.getOutputTimersFlatten().size());
        assertEquals(BASE_HIT_TIME + WAIT_SECONDS,
                response.getOutputTimersFlatten().get(0).getTriggerTimestamp());

        // Состояние должно содержать hit_payload
        Payload key = buildKey("hit-1", BASE_HIT_TIME);
        var state = response.getExternalState("join-state", key);
        assertFalse(state.isReset());
        assertEquals("payload-data", state.getValue().get("hit_payload", String.class));
    }

    @Test
    void testTimerEmitsJoinedAction() {
        Payload key = buildKey("hit-10", BASE_HIT_TIME);

        // Предзаполняем состояние
        ExternalState preState = new ExternalState(
                new PayloadBuilder(joinStateSchema)
                        .set("hit_payload", "some-payload")
                        .set("show_time", BASE_HIT_TIME + 3L)
                        .set("click_time", BASE_HIT_TIME + 7L)
                        .finish()
        );

        Timer timer = Timer.builder()
                .setStreamId("timer")
                .setKey(key)
                .setTriggerTimestamp(BASE_HIT_TIME + WAIT_SECONDS)
                .build();

        var request = TestDoProcessRequest.builder("join")
                .setTimers(List.of(timer))
                .setExternalState("join-state", Map.of(key, preState))
                .build();

        var response = testHarness.doProcess(request);

        // Должно быть одно выходное сообщение
        assertEquals(1, response.getOutputMessagesFlatten().size());
        var msg = response.getOutputMessagesFlatten().get(0);
        assertEquals("joined_action", msg.getStreamId());

        JoinedAction result = (JoinedAction) msg.getPayload();
        assertEquals("hit-10", result.getHitId());
        assertTrue(result.getClick());

        // Состояние должно быть сброшено
        assertTrue(response.getExternalState("join-state", key).isReset());
    }

    @Test
    void testLateMessageIsDropped() {
        long watermark = BASE_HIT_TIME + 5L;
        long eventTimestamp = BASE_HIT_TIME + 3L; // < watermark → late

        var messages = List.of(buildActionMessage("hit-5", BASE_HIT_TIME,
                BASE_HIT_TIME + 2L, false, eventTimestamp));

        var request = TestDoProcessRequest.builder("join")
                .setMessages(messages)
                .setWatermarks(Map.of("hit", watermark, "action", watermark))
                .build();

        var response = testHarness.doProcess(request);

        // Опоздавшее сообщение отброшено
        assertTrue(response.getOutputMessagesFlatten().isEmpty());
        assertTrue(response.getOutputTimersFlatten().isEmpty());
        assertEquals(0, response.getExternalStateSize("join-state"));
    }
}

@Test
void testFullJoinFlow() {
    String hitId = "hit-20";
    long hitTime = BASE_HIT_TIME;

    // Шаг 1: обработка hit-сообщения
    var hitRequest = TestDoProcessRequest.builder("join")
            .setMessages(List.of(buildHitMessage(hitId, hitTime, "full-payload")))
            .setWatermarks(defaultWatermarks())
            .build();
    var hitResponse = testHarness.doProcess(hitRequest);

    Payload key = buildKey(hitId, hitTime);
    var stateAfterHit = hitResponse.getExternalState("join-state", key);

    // Шаг 2: обработка show-сообщения, передаём состояние из шага 1
    var showRequest = TestDoProcessRequest.builder("join")
            .setMessages(List.of(buildActionMessage(hitId, hitTime, hitTime + 3L, false, hitTime + 1L)))
            .setExternalState("join-state", Map.of(key, stateAfterHit))
            .setWatermarks(defaultWatermarks())
            .build();
    var showResponse = testHarness.doProcess(showRequest);

    var stateAfterShow = showResponse.getExternalState("join-state", key);

    // Шаг 3: обработка click-сообщения
    var clickRequest = TestDoProcessRequest.builder("join")
            .setMessages(List.of(buildActionMessage(hitId, hitTime, hitTime + 7L, true, hitTime + 2L)))
            .setExternalState("join-state", Map.of(key, stateAfterShow))
            .setWatermarks(defaultWatermarks())
            .build();
    var clickResponse = testHarness.doProcess(clickRequest);

    var stateAfterClick = clickResponse.getExternalState("join-state", key);

    // Шаг 4: срабатывание таймера
    Timer timer = Timer.builder()
            .setStreamId("timer")
            .setKey(key)
            .setTriggerTimestamp(hitTime + WAIT_SECONDS)
            .build();
    var timerRequest = TestDoProcessRequest.builder("join")
            .setTimers(List.of(timer))
            .setExternalState("join-state", Map.of(key, stateAfterClick))
            .build();
    var timerResponse = testHarness.doProcess(timerRequest);

    // Проверяем финальный результат
    assertEquals(1, timerResponse.getOutputMessagesFlatten().size());
    JoinedAction result = timerResponse.getOutputMessagesFlatten().get(0).getPayload();
    assertEquals(hitId, result.getHitId());
    assertTrue(result.getClick());

    assertTrue(timerResponse.getExternalState("join-state", key).isReset());
}

@TestConfiguration
public class MyTestConfiguration {

    @Bean
    public CompanionExecutionConfig companionExecutionConfig() {
        return new CompanionExecutionConfig(0, new MockEnvironmentReader().worker());
    }

    @Bean
    public GrpcServerExecution grpcServerExecution(
            PipelineContext pipelineContext,
            CompanionExecutionConfig companionExecutionConfig
    ) {
        return new NoServerTestExecution(pipelineContext, companionExecutionConfig);
    }
}

@SpringBootTest(classes = MyTestConfiguration.class)
class JoinFunctionTest {

    @Autowired
    private PipelineContext pipelineContext;

    private TestComputationHarness testHarness;
    private TableSchema keySchema;

    @BeforeEach
    void init() throws IOException {
        this.keySchema = TableSchema.builder()
                .addValue("hash", ColumnValueType.UINT64)
                .addValue("hit_id", ColumnValueType.STRING)
                .addValue("hit_time", ColumnValueType.UINT64)
                .build();

        var specPath = Paths.getSourcePath(
                "yt/yt/flow/examples/java/lb_wait_click_join/test/pipeline.yson");
        var txtSpec = Files.readString(Path.of(specPath));

        this.testHarness = TestComputationHarness.builder()
                .setPipelineContext(pipelineContext)  // инжектированный через Spring
                .setPipelineSpec(txtSpec)
                .build();
    }

    @Test
    void testHitMessageStoresState() {
        var messages = List.of(buildHitMessage("hit-1", 1000L, "payload-data"));

        var request = TestDoProcessRequest.builder("join")
                .setMessages(messages)
                .setWatermarks(Map.of("hit", 0L, "action", 0L))
                .build();

        var response = testHarness.doProcess(request);

        assertTrue(response.getOutputMessagesFlatten().isEmpty());
        assertEquals(1, response.getOutputTimersFlatten().size());
        assertEquals(1, response.getInternalStateSize("join-state"));
    }

    // ... вспомогательные методы и остальные тесты
}