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

PEERDIR(
    yt/yt/flow/library/python/companion
)

from yt.yt.flow.library.python.companion.context import RequestContext
from yt.yt.flow.library.python.companion.job import Job
from yt.yt.flow.library.python.companion.stream import (
    StreamIdsMapping,
    StreamSpecs,
    RawStream,
)
from yt.yt.flow.library.python.companion.row import ExtendedMessage


def make_request_ctx(messages=None, timers=None):
    """Создание минимального RequestContext для тестов."""
    mapping = StreamIdsMapping({"input": 0, "output": 1})
    specs = StreamSpecs(mapping, [RawStream("input"), RawStream("output")])

    job = Job(
        job_id="test-job",
        computation_id="test-comp",
        stream_specs=specs,
        static_spec={},
    )

    return RequestContext(
        job_id="test-job",
        request_id="test-req",
        computation_id="test-comp",
        messages=messages or [],
        timers=timers or [],
        stream_specs=specs,
        job=job,
    )

from yt.yt.flow.library.python.companion.computation import (
    Computation,
    RowFunction,
)
from yt.yt.flow.library.python.companion.row import (
    ExtendedMessage,
    Message,
)


class PassthroughFunction(RowFunction):
    def on_message(self, message, output, ctx):
        output.add_message(
            Message(message_id=message.message_id, stream_id="output")
        )


def test_passthrough():
    comp = Computation(
        computation_id="test",
        process_function=PassthroughFunction(),
    )
    messages = [
        ExtendedMessage(message_id="m1", stream_id="input"),
        ExtendedMessage(message_id="m2", stream_id="input"),
    ]
    ctx = make_request_ctx(messages=messages)
    response = comp.do_process(ctx)

    assert len(response.transform_results) == 2
    assert response.transform_results[0].messages[0].message_id == "m1"
    assert response.transform_results[1].messages[0].message_id == "m2"

from yt.yt.flow.library.python.companion.row import Timer


class TimerFunction(RowFunction):
    def on_message(self, message, output, ctx):
        output.add_timer(trigger_timestamp=1000, event_timestamp=500)

    def on_timer(self, timer, output, ctx):
        output.add_message(
            Message(message_id="from-timer", stream_id="output")
        )


def test_timer_roundtrip():
    comp = Computation(
        computation_id="test",
        process_function=TimerFunction(),
    )
    messages = [ExtendedMessage(message_id="m1")]
    timers = [Timer(message_id="t1")]
    ctx = make_request_ctx(messages=messages, timers=timers)
    response = comp.do_process(ctx)

    results = response.transform_results
    # Сообщение создаёт таймер, таймер создаёт сообщение
    assert any(r.timers for r in results)
    assert any(r.messages for r in results)

import yt.type_info as ti

from yt.yt.flow.library.python.companion.context import DefaultRuntimeContext
from yt.yt.flow.library.python.companion.row import (
    ColumnSchema,
    ExtendedMessage,
    Payload,
    TableSchema,
)
from yt.yt.flow.library.python.companion.stream import (
    StreamIdsMapping,
    StreamSpecs,
    RawStream,
)
from yt.yt.flow.library.python.companion.wire_protocol import (
    ColumnValueType,
    UnversionedRow,
    UnversionedValue,
)


KEY_SCHEMA = TableSchema([ColumnSchema("id", ti.String)])


def make_key_payload():
    row = UnversionedRow(values=[
        UnversionedValue(column_id=0, type=ColumnValueType.STRING, value=b"test-key"),
    ])
    return Payload(row, KEY_SCHEMA)


def make_ctx(internal_state_names=None, **kwargs):
    streams = [RawStream("input"), RawStream("output")]
    mapping = StreamIdsMapping({s.stream_id: i for i, s in enumerate(streams)})
    specs = StreamSpecs(mapping, streams)
    return DefaultRuntimeContext(
        internal_state_names=internal_state_names or set(),
        stream_specs=specs,
        internal_states=kwargs.get("internal_states", {}),
        external_states=kwargs.get("external_states", {}),
        watermarks={},
        min_watermark=0,
        computation_parameters={},
        key_schema=KEY_SCHEMA,
    )


def test_yson_state_roundtrip():
    ctx = make_ctx(internal_state_names={"word-state"})
    message = ExtendedMessage(message_id="m1", key=make_key_payload())

    accessor = ctx.state("word-state", message)
    accessor.set({"word": "hello", "count": 1})

    accessor2 = ctx.state("word-state", message)
    result = accessor2.get()
    assert result["word"] == "hello"
    assert result["count"] == 1

from google.protobuf.wrappers_pb2 import Int64Value


def test_proto_state_roundtrip():
    ctx = make_ctx(internal_state_names={"join-state"})
    message = ExtendedMessage(message_id="m1", key=make_key_payload())

    accessor = ctx.proto_state("join-state", message, Int64Value)
    state = Int64Value(value=42)
    accessor.set(state)

    accessor2 = ctx.proto_state("join-state", message, Int64Value)
    assert accessor2.get().value == 42

from yt.yt.flow.library.python.companion.state import StatesHolder
from yt.yt.flow.library.python.companion.row import PayloadBuilder


STATE_SCHEMA = TableSchema([
    ColumnSchema("count", ti.Int64),
    ColumnSchema("name", ti.String),
])


def test_external_state_roundtrip():
    ext_holder = StatesHolder("ext", KEY_SCHEMA, STATE_SCHEMA)
    ctx = make_ctx(external_states={"/shuffle-state": ext_holder})
    message = ExtendedMessage(message_id="m1", key=make_key_payload())

    state = ctx.external_state("/shuffle-state", message)
    builder = state.to_builder()
    builder.set("count", 99)
    state.set(builder.finish())

    state2 = ctx.external_state("/shuffle-state", message)
    assert state2.get("count") == 99

PY3TEST()

INCLUDE(${ARCADIA_ROOT}/yt/yt/flow/library/python/integration_test_base/recipe.inc)

TEST_SRCS(
    test_wordcount.py
    yt_sync.py
)

PEERDIR(
    yt/yt/flow/library/python/queue
)

DEPENDS(
    ${MODDIR}/..
    yt/yt/flow/bin/flow_server
)

DATA(arcadia/${MODDIR}/pipeline.yson)

REQUIREMENTS(
    cpu:4
    ram:32
)

TAG(ya:huge_logs)

SIZE(MEDIUM)

END()

class Test(FlowTestPythonBase):
    PYTHON_COMPANION_BINARY = yatest.common.binary_path("yt/yt/flow/examples/python/word_count/word_count")

alias curtestdir="ps -f -u $USER | python3 -c \"import sys, re; drs = set(e for e in re.findall(r'[\s=](/\S*testing' + r'_out_stuff)\b', sys.stdin.read())); print('' if len(drs) == 1 else 'Select first from ' + repr(drs), file=sys.stderr); print(list(drs)[0])\""
alias cdcurtestdir='cd $(curtestdir)'

# Достать из логов ссылку на UI локального YTsaurus.
alias curlocalyt='cat $(curtestdir)/stderr 2>/dev/null | grep YT'

ya make -A --test-param RUNNER_LOG_LEVEL=Debug
ya make -A --test-disable-timeout --test-param PAUSE_BEFORE_FLOW_PROCESS_FEDERATION_TEARDOWN=1

ya make -A --test-param 'EXTERNAL_YT_CONFIG={path="//tmp/yt_flow";tablet_cell_bundle=default;proxy_role=default;clusters=[{cluster_name=<cluster-name>};];}'

ya make -A --ext-py