numaflow-python/packages/pynumaflow/examples/sink/log/test_example.py at 2b20b6790d838bb1e55373e15f7055d99ecc0c3a · numaproj/numaflow-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
"""Tests for the log sink example handlers."""

import pytest

from example import UserDefinedSink, udsink_handler
from pynumaflow.sinker.testing import SinkerTestKit, datum


@pytest.fixture(scope="module")
def sinker_kit():
    with SinkerTestKit(udsink_handler) as kit:
        yield kit


class TestFuncHandler:
    def test_is_ready(self, sinker_kit):
        with sinker_kit.client() as c:
            assert c.is_ready()

    def test_single_datum(self, sinker_kit):
        with sinker_kit.client() as c:
            result = c.send_request(datum(b"hello"))
        assert len(result) == 1
        result.assert_all_success()

    def test_multiple_datums(self, sinker_kit):
        with sinker_kit.client() as c:
            result = c.send_request(
                datum(b"msg-1"),
                datum(b"msg-2"),
                datum(b"msg-3"),
            )
        assert len(result) == 3
        result.assert_all_success()
        result.assert_counts(success=3, failed=0, fallback=0)

    def test_str_payload(self, sinker_kit):
        with sinker_kit.client() as c:
            result = c.send_request(datum("string payload"))
        result.assert_all_success()

    def test_dict_payload(self, sinker_kit):
        with sinker_kit.client() as c:
            result = c.send_request(datum({"event": "click", "count": 1}))
        result.assert_all_success()

    def test_with_keys(self, sinker_kit):
        with sinker_kit.client() as c:
            result = c.send_request(datum(b"keyed", keys=["k1", "k2"]))
        result.assert_all_success()

    def test_with_headers(self, sinker_kit):
        with sinker_kit.client() as c:
            result = c.send_request(
                datum(b"with-headers", headers={"x-trace-id": "abc123"})
            )
        result.assert_all_success()

    def test_auto_coerce(self, sinker_kit):
        """send_request accepts raw bytes/str/dict without wrapping in datum()."""
        with sinker_kit.client() as c:
            result = c.send_request(b"raw", "text", {"k": "v"})
        assert len(result) == 3
        result.assert_all_success()

    def test_large_batch(self, sinker_kit):
        with sinker_kit.client() as c:
            result = c.send_request(*[datum(f"item-{i}") for i in range(50)])
        assert len(result) == 50
        result.assert_all_success()


# ---------------------------------------------------------------------------
# Class-based handler (UserDefinedSink)
# ---------------------------------------------------------------------------

@pytest.fixture(scope="module")
def class_kit():
    with SinkerTestKit(UserDefinedSink()) as kit:
        yield kit


class TestClassHandler:
    def test_is_ready(self, class_kit):
        with class_kit.client() as c:
            assert c.is_ready()

    def test_single_datum(self, class_kit):
        with class_kit.client() as c:
            result = c.send_request(datum(b"hello"))
        assert len(result) == 1
        result.assert_all_success()

    def test_multiple_datums(self, class_kit):
        with class_kit.client() as c:
            result = c.send_request(
                datum(b"msg-1"),
                datum(b"msg-2"),
                datum(b"msg-3"),
            )
        assert len(result) == 3
        result.assert_all_success()
        result.assert_counts(success=3, failed=0, fallback=0)

    def test_str_payload(self, class_kit):
        with class_kit.client() as c:
            result = c.send_request(datum("string payload"))
        result.assert_all_success()

    def test_dict_payload(self, class_kit):
        with class_kit.client() as c:
            result = c.send_request(datum({"event": "click", "count": 1}))
        result.assert_all_success()

    def test_with_keys(self, class_kit):
        with class_kit.client() as c:
            result = c.send_request(datum(b"keyed", keys=["k1", "k2"]))
        result.assert_all_success()

    def test_auto_coerce(self, class_kit):
        with class_kit.client() as c:
            result = c.send_request(b"raw", "text", {"k": "v"})
        assert len(result) == 3
        result.assert_all_success()

    def test_large_batch(self, class_kit):
        with class_kit.client() as c:
            result = c.send_request(*[datum(f"item-{i}") for i in range(50)])
        assert len(result) == 50
        result.assert_all_success()