wehub-resource-sync
172 lines
3.1 KiB
Python
172 lines
3.1 KiB
Python
import asyncio
|
|
|
|
import pytest
|
|
|
|
from e2b_connect.client import (
|
|
EnvelopeFlags,
|
|
ServerStreamParser,
|
|
_retry,
|
|
encode_envelope,
|
|
)
|
|
|
|
|
|
class GoodError(Exception):
|
|
pass
|
|
|
|
|
|
class BadError(Exception):
|
|
pass
|
|
|
|
|
|
def test_sync_retry_after_expected_exception():
|
|
total = 0
|
|
|
|
@_retry(GoodError, 1)
|
|
def f():
|
|
nonlocal total
|
|
total += 1
|
|
raise GoodError()
|
|
|
|
with pytest.raises(GoodError):
|
|
f()
|
|
|
|
assert total == 2
|
|
|
|
|
|
def test_sync_do_not_retry_on_unexpected_exception():
|
|
total = 0
|
|
|
|
@_retry(GoodError, 1)
|
|
def f():
|
|
nonlocal total
|
|
total += 1
|
|
raise BadError()
|
|
|
|
with pytest.raises(BadError):
|
|
f()
|
|
|
|
assert total == 1
|
|
|
|
|
|
def test_sync_do_not_throw_when_retry_works():
|
|
total = 0
|
|
|
|
@_retry(GoodError, 1)
|
|
def f():
|
|
nonlocal total
|
|
total += 1
|
|
|
|
if total < 2:
|
|
raise GoodError()
|
|
|
|
return True
|
|
|
|
result = f()
|
|
assert result is True
|
|
assert total == 2
|
|
|
|
|
|
async def test_async_retry_after_expected_exception():
|
|
total = 0
|
|
|
|
@_retry(GoodError, 1)
|
|
async def f():
|
|
nonlocal total
|
|
total += 1
|
|
raise GoodError()
|
|
|
|
with pytest.raises(GoodError):
|
|
await f()
|
|
|
|
assert total == 2
|
|
|
|
|
|
async def test_async_do_not_retry_on_unexpected_exception():
|
|
total = 0
|
|
|
|
@_retry(GoodError, 1)
|
|
async def f():
|
|
nonlocal total
|
|
total += 1
|
|
raise BadError()
|
|
|
|
with pytest.raises(BadError):
|
|
await f()
|
|
|
|
assert total == 1
|
|
|
|
|
|
async def test_async_do_not_throw_when_retry_works():
|
|
total = 0
|
|
|
|
@_retry(GoodError, 1)
|
|
async def f():
|
|
nonlocal total
|
|
total += 1
|
|
|
|
if total < 2:
|
|
raise GoodError()
|
|
|
|
return True
|
|
|
|
result = await f()
|
|
assert result is True
|
|
assert total == 2
|
|
|
|
|
|
async def test_async_with_multiple_await_calls():
|
|
total = 0
|
|
|
|
async def a():
|
|
await asyncio.sleep(0.001)
|
|
|
|
@_retry(GoodError, 1)
|
|
async def f():
|
|
nonlocal total
|
|
total += 1
|
|
|
|
await a()
|
|
|
|
if total < 2:
|
|
raise GoodError()
|
|
|
|
await a()
|
|
|
|
return True
|
|
|
|
result = await f()
|
|
assert result is True
|
|
assert total == 2
|
|
|
|
|
|
def make_parser():
|
|
return ServerStreamParser(
|
|
decode=lambda data, msg_type: data,
|
|
response_type=None,
|
|
)
|
|
|
|
|
|
def test_parser_yields_messages_from_single_chunk():
|
|
parser = make_parser()
|
|
envelope = encode_envelope(flags=EnvelopeFlags(0), data=b"abc")
|
|
|
|
assert list(parser.parse(envelope)) == [b"abc"]
|
|
|
|
|
|
def test_parser_handles_payload_split_after_header():
|
|
parser = make_parser()
|
|
envelope = encode_envelope(flags=EnvelopeFlags(0), data=b"abc")
|
|
|
|
# The header plus one payload byte arrive first; the remaining payload
|
|
# (shorter than a header) arrives in a later chunk.
|
|
assert list(parser.parse(envelope[:6])) == []
|
|
assert list(parser.parse(envelope[6:])) == [b"abc"]
|
|
|
|
|
|
def test_parser_handles_end_stream_payload_shorter_than_header():
|
|
parser = make_parser()
|
|
envelope = encode_envelope(flags=EnvelopeFlags.end_stream, data=b"{}")
|
|
|
|
assert list(parser.parse(envelope[:5])) == []
|
|
assert list(parser.parse(envelope[5:])) == []
|