python-zyte-api/tests/test_sync.py at bad7988be67726f5164ee960568706ce17f471a2 · zytedata/python-zyte-api · 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
154
155
156
157
158
159
160
161
162
from __future__ import annotations

from types import GeneratorType
from typing import TYPE_CHECKING, Any
from unittest.mock import AsyncMock, patch

import pytest

from zyte_api import ZyteAPI
from zyte_api.apikey import NoApiKey

if TYPE_CHECKING:
    from tests.mockserver import MockServer


def test_api_key():
    ZyteAPI(api_key="a")
    with pytest.raises(NoApiKey):
        ZyteAPI()


def test_trust_env_is_forwarded():
    with patch("zyte_api._sync.AsyncZyteAPI") as async_client:
        ZyteAPI(api_key="a", trust_env=True)
    assert async_client.call_args.kwargs["trust_env"] is True


def test_get(mockserver):
    client = ZyteAPI(api_key="a", api_url=mockserver.urljoin("/"))
    expected_result = {
        "url": "https://a.example",
        "httpResponseBody": "PGh0bWw+PGJvZHk+SGVsbG88aDE+V29ybGQhPC9oMT48L2JvZHk+PC9odG1sPg==",
    }
    actual_result = client.get({"url": "https://a.example", "httpResponseBody": True})
    assert actual_result == expected_result


def test_iter(mockserver):
    client = ZyteAPI(api_key="a", api_url=mockserver.urljoin("/"))
    queries = [
        {"url": "https://a.example", "httpResponseBody": True},
        {"url": "https://exception.example", "httpResponseBody": True},
        {"url": "https://b.example", "httpResponseBody": True},
    ]
    expected_results = [
        {
            "url": "https://a.example",
            "httpResponseBody": "PGh0bWw+PGJvZHk+SGVsbG88aDE+V29ybGQhPC9oMT48L2JvZHk+PC9odG1sPg==",
        },
        Exception,
        {
            "url": "https://b.example",
            "httpResponseBody": "PGh0bWw+PGJvZHk+SGVsbG88aDE+V29ybGQhPC9oMT48L2JvZHk+PC9odG1sPg==",
        },
    ]
    actual_results = client.iter(queries)
    assert isinstance(actual_results, GeneratorType)
    actual_results_list = list(actual_results)
    assert len(actual_results_list) == len(expected_results)
    for actual_result in actual_results_list:
        if isinstance(actual_result, Exception):
            assert Exception in expected_results
        else:
            assert actual_result in expected_results


def test_semaphore(mockserver):
    client = ZyteAPI(api_key="a", api_url=mockserver.urljoin("/"))
    client._async_client._semaphore = AsyncMock(wraps=client._async_client._semaphore)
    queries = [
        {"url": "https://a.example", "httpResponseBody": True},
        {"url": "https://b.example", "httpResponseBody": True},
        {"url": "https://c.example", "httpResponseBody": True},
    ]
    client.get(queries[0])
    next(iter(client.iter(queries[1:2])))
    client.get(queries[2])
    assert client._async_client._semaphore.__aenter__.call_count == len(queries)
    assert client._async_client._semaphore.__aexit__.call_count == len(queries)


def test_session_context_manager(mockserver: MockServer) -> None:
    client = ZyteAPI(api_key="a", api_url=mockserver.urljoin("/"))
    queries = [
        {"url": "https://a.example", "httpResponseBody": True},
        {"url": "https://exception.example", "httpResponseBody": True},
        {"url": "https://b.example", "httpResponseBody": True},
    ]
    expected_results = [
        {
            "url": "https://a.example",
            "httpResponseBody": "PGh0bWw+PGJvZHk+SGVsbG88aDE+V29ybGQhPC9oMT48L2JvZHk+PC9odG1sPg==",
        },
        Exception,
        {
            "url": "https://b.example",
            "httpResponseBody": "PGh0bWw+PGJvZHk+SGVsbG88aDE+V29ybGQhPC9oMT48L2JvZHk+PC9odG1sPg==",
        },
    ]
    actual_results: list[dict[str, Any] | Exception] = []
    with client.session() as session:
        assert session._session.connector is not None
        assert session._session.connector.limit == client._async_client.n_conn
        actual_results.append(session.get(queries[0]))
        actual_results.extend(session.iter(queries[1:]))
        aiohttp_session = session._session
        assert not aiohttp_session.closed
    assert aiohttp_session.closed

    with pytest.raises(RuntimeError):
        session.get(queries[0])

    assert isinstance(next(iter(session.iter(queries[1:]))), RuntimeError)

    assert len(actual_results) == len(expected_results)
    for actual_result in actual_results:
        if isinstance(actual_result, Exception):
            assert Exception in expected_results
        else:
            assert actual_result in expected_results


def test_session_no_context_manager(mockserver: MockServer) -> None:
    client = ZyteAPI(api_key="a", api_url=mockserver.urljoin("/"))
    queries = [
        {"url": "https://a.example", "httpResponseBody": True},
        {"url": "https://exception.example", "httpResponseBody": True},
        {"url": "https://b.example", "httpResponseBody": True},
    ]
    expected_results = [
        {
            "url": "https://a.example",
            "httpResponseBody": "PGh0bWw+PGJvZHk+SGVsbG88aDE+V29ybGQhPC9oMT48L2JvZHk+PC9odG1sPg==",
        },
        Exception,
        {
            "url": "https://b.example",
            "httpResponseBody": "PGh0bWw+PGJvZHk+SGVsbG88aDE+V29ybGQhPC9oMT48L2JvZHk+PC9odG1sPg==",
        },
    ]
    actual_results: list[dict[str, Any] | Exception] = []
    session = client.session()
    assert session._session.connector is not None
    assert session._session.connector.limit == client._async_client.n_conn
    actual_results.append(session.get(queries[0]))
    actual_results.extend(session.iter(queries[1:]))
    aiohttp_session = session._session
    assert not aiohttp_session.closed
    session.close()
    assert aiohttp_session.closed

    with pytest.raises(RuntimeError):
        session.get(queries[0])

    assert isinstance(next(iter(session.iter(queries[1:]))), RuntimeError)

    assert len(actual_results) == len(expected_results)
    for actual_result in actual_results:
        if isinstance(actual_result, Exception):
            assert Exception in expected_results
        else:
            assert actual_result in expected_results