dream

1
import allure
2
import pytest
3
import requests
4

5
from typing import Dict, List
6

7

8
@allure.description("""Test entities detection and labeling""")
9
@pytest.mark.parametrize(
10
    "request_data, gold_results",
11
    [
12
        (
13
            {"sentences": [["what is the capital of russia?"]]},
14
            [
15
                {
16
                    "entities": ["capital", "russia"],
17
                    "labelled_entities": [
18
                        {
19
                            "finegrained_label": [["misc", 0.871]],
20
                            "label": "misc",
21
                            "offsets": [12, 19],
22
                            "text": "capital",
23
                        },
24
                        {
25
                            "finegrained_label": [["loc", 0.9927]],
26
                            "label": "location",
27
                            "offsets": [23, 29],
28
                            "text": "russia",
29
                        },
30
                    ],
31
                }
32
            ],
33
        ),
34
        (
35
            {"sentences": [["let's talk about politics."]]},
36
            [
37
                {
38
                    "entities": ["politics"],
39
                    "labelled_entities": [
40
                        {
41
                            "finegrained_label": [["misc", 0.9984]],
42
                            "label": "misc",
43
                            "offsets": [17, 25],
44
                            "text": "politics",
45
                        }
46
                    ],
47
                }
48
            ],
49
        ),
50
    ],
51
)
52
def test_entity_detection(url: str, request_data: Dict[str, list], gold_results: List[Dict]):
53
    response = requests.post(url, json=request_data)
54
    result = response.json()
55
    assert response.status_code == 200
56
    assert "entities" in result[0]
57
    assert "labelled_entities" in result[0]
58
    assert result == gold_results
59