1
# Copyright (c) 2023 Cloudflare, Inc.
2
# Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
#     https://opensource.org/licenses/Apache-2.0
4
 
5
import collections.abc
6
import json
7
from asyncio import Future, sleep
8
from datetime import datetime
9
from http import HTTPMethod
10
from unittest import TestCase
11
 
12
import js
13
from workers import Blob, Request, Response, WorkerEntrypoint, handler
14
 
15
from pyodide.ffi import JsException, JsProxy, to_js
16
 
17
assertRaises = TestCase().assertRaises
18
assertRaisesRegex = TestCase().assertRaisesRegex
19
 
20
testFuture = Future()
21
 
22
 
23
class PythonRpcTester(WorkerEntrypoint):
24
    def __init__(self, ctx, env):
25
        super().__init__(ctx, env)
26
        # Verify that the superclass constructor initialises the env/ctx fields.
27
        assert self.env is not None
28
        assert self.ctx is not None
29
 
30
        assert not isinstance(env, JsProxy)
31
        self.env = env
32
 
33
    async def no_args(self):
34
        return "hello from python"
35
 
36
    async def one_arg(self, x):
37
        return f"{x}"
38
 
39
    async def identity(self, x):
40
        assert not isinstance(x, JsProxy)
41
        return x
42
 
43
    async def handle_response(self, response):
44
        # Verify that we receive a Python object here...
45
        assert isinstance(response, Response)
46
        return response
47
 
48
    async def handle_request(self, req):
49
        assert isinstance(req, Request)
50
        return req
51
 
52
    async def check_env(self):
53
        # Verify that the `env` supplied to the entrypoint class is wrapped.
54
        curr_date = datetime.now()
55
        py_date = await self.env.PythonRpc.identity(curr_date)
56
        assert isinstance(py_date, datetime)
57
        # JavaScript Date objects only have millisecond precision
58
        assert abs((py_date - curr_date).microseconds) < 1e6
59
 
60
        return True
61
 
62
    async def sleep_then_set_result(self):
63
        await sleep(0.1)
64
        testFuture.set_result(100)
65
 
66
    async def test_wait_until_coroutine_lifetime(self):
67
        self.ctx.waitUntil(self.sleep_then_set_result())
68
 
69
 
70
class CustomType:
71
    def __init__(self, x):
72
        self.x = 42
73
 
74
 
75
def assert_equal(a, b, avoid_type_check):
76
    # Convert to JSON so that the contents of the values can be easily compared.
77
    received = json.dumps(
78
        json.loads(js.JSON.stringify(a) if isinstance(a, JsProxy) else json.dumps(a))
79
    )
80
    expected = json.dumps(
81
        json.loads(js.JSON.stringify(b) if isinstance(b, JsProxy) else json.dumps(b))
82
    )
83
    if received != expected:
84
        raise ValueError(
85
            f"Assert failed, args contents are not equal. received='{received}' expected='{expected}'"
86
        )
87
 
88
    if type(a) is not type(b) and not avoid_type_check:
89
        raise ValueError(
90
            f"Assert failed, types don't match. received={type(a)} expected={type(b)}"
91
        )
92
 
93
 
94
@handler
95
async def test(ctrl, env, ctx):
96
    # https://developer.mozilla.org/en-US/docs/Web/API/Web_Workers_API/Structured_clone_algorithm#supported_types
97
    #
98
    # Workers RPC doesn't support all of the above, but does support all the JS Types and
99
    # some of the Web API types.
100
    #
101
    # ReadableStream and WritableStream are also supported.
102
    #
103
    # We verify that we can send and receive as much of these as possible, including between
104
    # Python<->Python, JS<->Python and vice versa. We also ensure that the types we receive are
105
    # native Python types, rather than `JsProxy`s.
106
 
107
    # Simple tests.
108
    assert await env.PythonRpc.no_args() == "hello from python"
109
    assert await env.JsRpc.noArgs() == "hello from js"
110
    assert await env.PythonRpc.one_arg(42) == "42"
111
    assert await env.JsRpc.oneArg(42) == "42"
112
    arr = await env.PythonRpc.identity([1, 2, 3])
113
    assert isinstance(arr, collections.abc.Sequence) and not isinstance(arr, str)
114
 
115
    # Verify that text bindings can be accessed.
116
    assert env.FOO == "text binding"
117
 
118
    # Python Types
119
    for val in ["test", [1, 2, 3], {"key": 42}, 42, 1.2345, False, True, None]:
120
        received = await env.PythonRpc.identity(val)
121
        assert not isinstance(received, JsProxy), (
122
            "Expected the returned value from RPC to be a Python type."
123
        )
124
        assert_equal(received, val, False)
125
        received = await env.JsRpc.identity(val)
126
        assert not isinstance(received, JsProxy), (
127
            "Expected the returned value from RPC to be a Python type."
128
        )
129
        assert_equal(received, val, False)
130
 
131
    curr_date = datetime.now()
132
    py_date = await env.PythonRpc.identity(curr_date)
133
    assert isinstance(py_date, datetime)
134
    # JavaScript Date objects only have millisecond precision
135
    assert abs((py_date - curr_date).microseconds) < 1e6
136
 
137
    py_date_list = await env.PythonRpc.identity([datetime.now(), datetime.now()])
138
    for d in py_date_list:
139
        assert isinstance(d, datetime)
140
 
141
    py_set = await env.PythonRpc.identity({1, 2, 3})
142
    assert isinstance(py_set, set)
143
    assert 2 in py_set
144
    assert 42 not in py_set
145
 
146
    py_binary = await env.PythonRpc.identity(b"binary")
147
    assert isinstance(py_binary, memoryview)
148
    py_binary = await env.PythonRpc.identity(memoryview(b"abcefg"))
149
    assert isinstance(py_binary, memoryview)
150
 
151
    # JS types
152
    for val in [
153
        to_js([1, 2, 3, 4]),
154
        js.Number.new("1234"),
155
    ]:
156
        received = await env.PythonRpc.identity(val)
157
        assert not isinstance(received, JsProxy), (
158
            "Expected the returned value from RPC to be a Python type."
159
        )
160
        assert_equal(received, val.to_py(), True)
161
        received = await env.JsRpc.identity(val)
162
        assert not isinstance(received, JsProxy), (
163
            "Expected the returned value from RPC to be a Python type."
164
        )
165
        assert_equal(received, val.to_py(), True)
166
 
167
    for val in [
168
        js.ArrayBuffer.new(8),
169
        js.DataView.new(js.ArrayBuffer.new(16)),
170
        js.Int16Array.of("10", "24"),
171
    ]:
172
        js_binary = await env.PythonRpc.identity(val)
173
        assert isinstance(js_binary, memoryview)
174
        js_binary = await env.JsRpc.identity(val)
175
        assert isinstance(js_binary, memoryview)
176
 
177
    js_map = await env.PythonRpc.identity(
178
        js.Map.new(
179
            [
180
                [1, "one"],
181
                [2, "two"],
182
                [3, "three"],
183
            ]
184
        ),
185
    )
186
    assert isinstance(js_map, dict)
187
    assert js_map[1] == "one"
188
 
189
    js_set = await env.PythonRpc.identity(
190
        js.Set.new([1, 2, 3, 4]),
191
    )
192
    assert isinstance(js_set, set)
193
    assert 1 in js_set
194
    assert 42 not in js_set
195
 
196
    js_undefined = await env.PythonRpc.identity(js.undefined)
197
    assert js_undefined is None
198
 
199
    js_date = await env.PythonRpc.identity(js.Date.new())
200
    assert isinstance(js_date, datetime)
201
 
202
    js_date_list = await env.PythonRpc.identity([js.Date.new(), js.Date.new()])
203
    for d in js_date_list:
204
        assert isinstance(d, datetime)
205
 
206
    js_exception = await env.PythonRpc.identity(js.Error.new("message"))
207
    assert isinstance(js_exception, Exception)
208
 
209
    js_obj = await env.PythonRpc.identity(
210
        to_js({"foo": 42}, dict_converter=js.Object.fromEntries)
211
    )
212
    assert isinstance(js_obj, dict)
213
    assert js_obj["foo"] == 42
214
 
215
    # Web/API Types
216
    # - Response
217
    py_response = await env.PythonRpc.handle_response(Response("this is a response"))
218
    assert isinstance(py_response, Response)
219
    assert await py_response.text() == "this is a response"
220
    js_response = await env.JsRpc.handleResponse(Response("this is a response"))
221
    assert await js_response.text() == "this is a response"
222
    assert isinstance(js_response, Response)
223
 
224
    # - Request
225
    py_response = await env.PythonRpc.handle_request(
226
        Request("https://test.com", method=HTTPMethod.POST)
227
    )
228
    assert isinstance(py_response, Request)
229
    assert py_response.method == "POST"
230
    js_response = await env.JsRpc.handleRequest(
231
        Request("https://test.com", method=HTTPMethod.POST)
232
    )
233
    assert js_response.method == "POST"
234
    assert isinstance(js_response, Request)
235
 
236
    # - Verify that a JS type can be sent.
237
    py_response2 = await env.PythonRpc.handle_response(js.Response.new("a JS response"))
238
    assert await py_response2.text() == "a JS response"
239
 
240
    # Verify that sending unsupported types fails.
241
    data_clone_regex = "^DataCloneError"
242
    with assertRaisesRegex(JsException, data_clone_regex):
243
        await env.PythonRpc.one_arg(CustomType(42))
244
    with assertRaisesRegex(JsException, data_clone_regex):
245
        await env.PythonRpc.one_arg(Blob("print(42)", content_type="text/python"))
246
    with assertRaisesRegex(JsException, data_clone_regex):
247
        await env.PythonRpc.identity(complex(1.23))
248
    with assertRaises(TypeError):
249
        await env.PythonRpc.identity((1, 2, 3))
250
    with assertRaisesRegex(JsException, data_clone_regex):
251
        await env.PythonRpc.identity(range(0, 30, 5))
252
    with assertRaises(TypeError):
253
        await env.PythonRpc.identity(bytearray.fromhex("2Ef0 F1f2  "))
254
    with assertRaises(TypeError):
255
        await env.PythonRpc.identity([(1, 2, 3)])
256
    # TODO: Support RegExp.
257
    with assertRaises(TypeError):
258
        await env.PythonRpc.identity(js.RegExp.new("ab+c", "i"))
259
    with assertRaises(TypeError):
260
        await env.PythonRpc.identity(lambda x: x + x)
261
    with assertRaises(TypeError):
262
 
263
        def my_func():
264
            pass
265
 
266
        await env.PythonRpc.identity(my_func)
267
    with assertRaises(TypeError):
268
        await env.PythonRpc.identity({"test": (1, 2, 3)})
269
 
270
    # Verify that the `env` in the DO is correctly wrapped.
271
    assert await env.PythonRpc.check_env()
272
 
273
    # Check that the coroutine returned by sleep_then_set_result() lasts long enough.
274
    # sleep_then_set_result() resolves testFuture after sleeping for 100ms.
275
    await env.PythonRpc.test_wait_until_coroutine_lifetime()
276
    assert not testFuture.done()
277
    await sleep(0.2)
278
    assert testFuture.result() == 100