chatsky/tests/core/test_context.py at a2823df194cb15722415811e02e500cd3fdf989b · deeppavlov/chatsky · 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
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
from copy import copy
import pytest

from chatsky.core.context import Context
from chatsky.core.ctx_utils import ContextError
from chatsky.core.node_label import AbsoluteNodeLabel
from chatsky.core.message import Message, MessageInitTypes
from chatsky.core.script_function import BaseResponse, BaseProcessing
from chatsky.core.pipeline import Pipeline
from chatsky.core import RESPONSE, PRE_TRANSITION, PRE_RESPONSE


class TestLabels:
    @pytest.fixture
    def ctx(self, context_factory):
        return context_factory(forbidden_fields=["requests", "responses"])

    def test_raises_on_empty_labels(self, ctx: Context):
        with pytest.raises(ContextError):
            ctx.last_label

    def test_existing_labels(self, ctx: Context):
        ctx.labels[5] = ("flow", "node1")

        assert ctx.last_label == AbsoluteNodeLabel(flow_name="flow", node_name="node1")
        ctx.labels[6] = ("flow", "node2")
        assert ctx.labels.keys() == [5, 6]
        assert ctx.last_label == AbsoluteNodeLabel(flow_name="flow", node_name="node2")


class TestRequests:
    @pytest.fixture
    def ctx(self, context_factory):
        return context_factory(forbidden_fields=["labels", "responses"])

    def test_existing_requests(self, ctx: Context):
        ctx.requests[5] = Message(text="text1")
        assert ctx.last_request == Message(text="text1")
        ctx.requests[6] = "text2"
        assert ctx.requests.keys() == [5, 6]
        assert ctx.last_request == Message(text="text2")

    def test_empty_requests(self, ctx: Context):
        with pytest.raises(ContextError):
            ctx.last_request

        ctx.requests[1] = "text"
        assert ctx.last_request == Message(text="text")
        assert ctx.requests.keys() == [1]


class TestResponses:
    @pytest.fixture
    def ctx(self, context_factory):
        return context_factory(forbidden_fields=["labels", "requests"])

    def test_existing_responses(self, ctx: Context):
        ctx.responses[5] = Message(text="text1")
        assert ctx.last_response == Message(text="text1")
        ctx.responses[6] = "text2"
        assert ctx.responses.keys() == [5, 6]
        assert ctx.last_response == Message(text="text2")

    def test_empty_responses(self, ctx: Context):
        assert ctx.last_response is None

        ctx.responses[1] = "text"
        assert ctx.last_response == Message(text="text")
        assert ctx.responses.keys() == [1]


class TestTurns:
    @pytest.fixture
    def ctx(self, context_factory):
        return context_factory()

    async def test_negative_index(self, ctx: Context):
        ctx.labels[5] = ("flow", "node5")
        ctx.requests[5] = Message(text="text5")
        ctx.responses[5] = Message(text="text5")
        ctx.current_turn_id = 5

        request, label, response = await ctx.turns[-1]
        assert label == AbsoluteNodeLabel(flow_name="flow", node_name="node5")
        assert request == Message(text="text5")
        assert response == Message(text="text5")

    async def test_partial_turn(self, ctx: Context):
        ctx.labels[6] = ("flow", "node6")
        ctx.requests[6] = Message(text="text6")
        ctx.current_turn_id = 6

        request, label, response = await ctx.turns[6]
        assert label == AbsoluteNodeLabel(flow_name="flow", node_name="node6")
        assert request == Message(text="text6")
        assert response is None

    async def test_slice_turn(self, ctx: Context):
        for i in range(2, 6):
            ctx.labels[i] = ("flow", f"node{i}")
            ctx.requests[i] = Message(text=f"text{i}")
            ctx.responses[i] = Message(text=f"text{i}")
            ctx.current_turn_id = i

        for i, turn in zip(range(2, 6), await ctx.turns[2:6]):
            request, label, response = turn
            assert AbsoluteNodeLabel(flow_name="flow", node_name=f"node{i}") == label
            assert Message(text=f"text{i}") == request
            assert Message(text=f"text{i}") == response


async def test_copy(context_factory):
    ctx = context_factory()
    ctx.misc["key"] = "value"

    cpy = copy(ctx)
    assert cpy.misc["key"] == "value"
    assert cpy._storage == ctx._storage
    assert cpy == ctx


async def test_pipeline_available():
    class MyResponse(BaseResponse):
        async def call(self, ctx: Context) -> MessageInitTypes:
            return ctx.pipeline.start_label.node_name

    pipeline = Pipeline(script={"flow": {"node": {RESPONSE: MyResponse()}}}, start_label=("flow", "node"))
    ctx = await pipeline._run_pipeline(Message(text=""), ctx_id="0")

    assert ctx.last_response == Message(text="node")

    ctx.framework_data.pipeline = None
    with pytest.raises(ContextError):
        await MyResponse().call(ctx)


async def test_current_node_available():
    log = []

    class MyProcessing(BaseProcessing):
        async def call(self, ctx: Context) -> None:
            log.append(ctx.current_node)

    pipeline = Pipeline(
        script={"flow": {"node": {PRE_RESPONSE: {"": MyProcessing()}, PRE_TRANSITION: {"": MyProcessing()}}}},
        start_label=("flow", "node"),
    )
    ctx = await pipeline._run_pipeline(Message(text=""), ctx_id="0")
    assert len(log) == 2

    ctx.framework_data.current_node = None
    with pytest.raises(ContextError):
        await MyProcessing().call(ctx)