aboutsummaryrefslogtreecommitdiffstatshomepage
path: root/tests/test_interceptor.py
blob: df5b675014ad3df61671d0962e91d80ac539beac (plain) (blame)
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
from infini.input import Input
from infini.interceptor import Interceptor
from infini.output import Output
from infini.router import ContainsRouter


def test_interceptor():
    input = Input("这个人叫简律纯.")
    valid_input = Input("这个叫苏向夜.")

    def intercept(_: Input) -> Input | Output:
        output = Output()
        output.block = True  # TODO 拦截器阻塞标识
        output.name = "block.jianlvchun"
        output.status = 0
        output.type = "text"
        return output

    interceptor = Interceptor()
    interceptor.interceptors = [
        {
            "priority": 1,
            "router": ContainsRouter("简律纯"),
            "handler": intercept,
        }
    ]
    output = interceptor.input(input)
    assert isinstance(output, Output)
    assert output.name == "block.jianlvchun"

    valid_output = interceptor.input(valid_input)
    assert isinstance(valid_output, Input)
    assert valid_output.get_plain_text() == "这个叫苏向夜."