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
from asyncio import sleep
6
from urllib.parse import urlparse
7
 
8
from js import Date
9
from workers import DurableObject, FetchResponse, Request, Response, WorkerEntrypoint
10
 
11
import pyodide
12
 
13
 
14
class MixinTest:
15
    def test_mixin(self):
16
        return 1234
17
 
18
 
19
class DurableObjectExample(DurableObject, MixinTest):
20
    def __init__(self, state, env):
21
        super().__init__(state, env)
22
        assert self.env is not None
23
        assert self.ctx is not None
24
 
25
        self.state = state
26
        self.counter = 0
27
        self.storage = state.storage
28
        self.alarm_triggered = False
29
 
30
        # Test blockConcurrencyWhile in the constructor with a Python async callback.
31
        # This is a common pattern for initializing DO state.
32
        async def init_callback():
33
            await self.storage.put("initialized", True)
34
 
35
        self.ctx.blockConcurrencyWhile(init_callback)
36
 
37
    async def fetch(self, request):
38
        assert isinstance(request, Request)
39
 
40
        curr = await self.storage.getAlarm()
41
        if not curr:
42
            self.storage.setAlarm(Date.now() + 100)
43
 
44
        url = urlparse(request.url)
45
        if url.path == "/counter":
46
            self.counter += 1
47
            return Response(f"hello from python {self.counter}")
48
        elif url.path == "/alarm":
49
            return Response(str(self.alarm_triggered))
50
        else:
51
            return Response("404")
52
 
53
    async def alarm(self, alarm_info):
54
        self.alarm_triggered = True
55
 
56
    async def no_args_method(self):
57
        return "value from python"
58
 
59
    async def args_method(self, arg):
60
        return "value from python " + arg
61
 
62
    def mutate_dict(self, my_dict):
63
        my_dict["foo"] = 42
64
 
65
    async def test_block_concurrency_while(self):
66
        # Verify the constructor's blockConcurrencyWhile ran successfully
67
        initialized = await self.storage.get("initialized")
68
        assert initialized, f"Expected True but got {initialized}"
69
 
70
        # Test blockConcurrencyWhile with a Python async callback that returns a value.
71
        async def my_callback():
72
            await self.storage.put("blocked", "yes")
73
            return 42
74
 
75
        result = await self.ctx.blockConcurrencyWhile(my_callback)
76
        assert result == 42, f"Expected 42 but got {result}"
77
        blocked = await self.storage.get("blocked")
78
        assert blocked == "yes", f"Expected 'yes' but got {blocked}"
79
 
80
        return True
81
 
82
    async def test_self_call(self):
83
        test_dict = dict()
84
        test_dict["test"] = 1
85
        self.mutate_dict(test_dict)
86
        assert test_dict["test"] == 1
87
        assert test_dict["foo"] == 42
88
        return True
89
 
90
    def jspi_method(self, arg):
91
        from pyodide.ffi import run_sync
92
 
93
        run_sync(sleep(0.01))
94
        return arg + 1
95
 
96
 
97
class Default(WorkerEntrypoint):
98
    async def test(self, ctrl):
99
        id = self.env.ns.idFromName("A")
100
        obj = self.env.ns.get(id)
101
 
102
        first_resp = await obj.fetch("http://foo.com/counter")
103
        first_resp_data = await first_resp.text()
104
        assert first_resp_data == "hello from python 1"
105
 
106
        second_resp = await obj.fetch("http://foo.com/counter")
107
        second_resp_data = await second_resp.text()
108
        assert second_resp_data == "hello from python 2"
109
 
110
        no_arg_resp = await obj.no_args_method()
111
        assert no_arg_resp == "value from python"
112
 
113
        arg_resp = await obj.args_method("test")
114
        assert arg_resp == "value from python test"
115
 
116
        assert await obj.test_block_concurrency_while()
117
 
118
        assert await obj.test_self_call()
119
 
120
        if pyodide.__version__ != "0.26.0a2":
121
            res = await obj.jspi_method(9)
122
            assert res == 10
123
 
124
        # Verify that a mixin method can be called via RPC.
125
        assert await obj.test_mixin() == 1234
126
 
127
        # Verify that DO fetch is wrapped.
128
        third_resp = await obj.fetch("http://foo.com/counter")
129
        assert isinstance(third_resp, FetchResponse)
130
        third_resp_data = await third_resp.text()
131
        assert third_resp_data == "hello from python 3"
132
 
133
        # Wait for alarm to get triggered.
134
        for _ in range(20):
135
            await sleep(0.2)
136
            resp = await obj.fetch("http://foo.com/alarm")
137
 
138
            alarm_triggered = await resp.text() == "True"
139
            if alarm_triggered:
140
                break
141
        else:
142
            raise AssertionError("Alarm never triggered")