def test_wait_for_in_executor(executor_class):
results = []
async def coro(func, loop, item, executor):
nonlocal results
results.append(await loop.run_in_executor(executor, func, item))
with entrypoint() as loop:
with executor_class(loop) as exec:
with pytest.raises(AssertionError):
loop.run_until_complete(
wait_for(*[
coro(blocking_bad_func, loop, i, exec)
for i in range(10)
]))
loop.run_until_complete(
wait_for(
*[coro(blocking_func, loop, i, exec) for i in range(10)]))
loop.run_until_complete(asyncio.sleep(1))
results.sort()
assert results
def test_wait_for_dummy():
with entrypoint() as loop:
results = loop.run_until_complete(
wait_for(*[asyncio.sleep(0.1) for _ in range(100)]))
assert len(results) == 100
assert results == [None] * 100
def loop(
request,
services,
loop_debug,
default_context,
entrypoint_kwargs,
thread_pool_size,
thread_pool_executor,
loop: asyncio.AbstractEventLoop,
caplog: pytest.LogCaptureFixture,
):
from aiomisc.context import get_context
from aiomisc.entrypoint import entrypoint
if LOG_LEVEL:
LOG_LEVEL.set(logging.getLogger().getEffectiveLevel())
pool = thread_pool_executor(thread_pool_size)
loop.set_default_executor(pool)
get_marker = request.node.get_closest_marker
forbid_loop_getter_marker = get_marker("forbid_get_event_loop")
catch_unhandled_marker = get_marker("catch_loop_exceptions")
exceptions = list()
if catch_unhandled_marker:
loop.set_exception_handler(lambda l, c: exceptions.append(c))
try:
with entrypoint(*services, loop=loop, **entrypoint_kwargs):
ctx = get_context(loop)
for key, value in default_context.items():
ctx[key] = value
if forbid_loop_getter_marker:
asyncio.get_event_loop.side_effect = partial(
pytest.fail,
"get_event_loop is forbidden",
)
yield loop
if exceptions:
logging.error(
"Unhandled exceptions found:\n\n\t%s",
"\n\t".join(("Message: {m}\n\t"
"Future: {f}\n\t"
"Exception: {e}").format(
m=e["message"],
f=repr(e.get("future")),
e=repr(e.get("exception")),
) for e in exceptions),
)
pytest.fail("Unhandled exceptions found. See logs.")
finally:
asyncio.get_event_loop.side_effect = get_event_loop
del loop
def test_simple():
class StartingService(Service):
async def start(self):
self.running = True
class DummyService(StartingService):
async def stop(self, err: Exception = None):
self.stopped = True
services = (
DummyService(running=False, stopped=False),
DummyService(running=False, stopped=False),
)
with entrypoint(*services):
pass
for svc in services:
assert svc.running
assert svc.stopped
services = (
DummyService(running=False, stopped=False),
DummyService(running=False, stopped=False),
)
with pytest.raises(RuntimeError):
with entrypoint(*services):
raise RuntimeError
for svc in services:
assert svc.running
assert svc.stopped
services = (
StartingService(running=False),
StartingService(running=False),
)
with pytest.raises(RuntimeError):
with entrypoint(*services):
raise RuntimeError
for svc in services:
assert svc.running
def run():
setproctitle(os.path.basename("[Worker] %s" % sys.argv[0]))
with entrypoint(*services,
pool_size=arguments.pool_size,
log_level=arguments.log_level,
log_format=arguments.log_format) as loop:
loop.set_debug(arguments.debug)
loop.run_forever()
def test_service_class():
with pytest.raises(NotImplementedError):
services = (
Service(running=False, stopped=False),
Service(running=False, stopped=False),
)
with entrypoint(*services):
pass
def test_aiohttp_service(unused_tcp_port):
@threaded
def http_client():
url = 'http://127.0.0.1:%s/' % unused_tcp_port
return requests.get(url, timeout=1).status_code
service = AIOHTTPTestApp(address='127.0.0.1', port=unused_tcp_port)
with entrypoint(service) as loop:
response = loop.run_until_complete(http_client())
assert response == 404
def loop(request, services, loop_debug, default_context, entrypoint_kwargs,
thread_pool_size, thread_pool_executor, loop):
from aiomisc.context import get_context
from aiomisc.entrypoint import entrypoint
asyncio.set_event_loop(loop)
pool = thread_pool_executor(thread_pool_size)
loop.set_default_executor(pool)
get_marker = request.node.get_closest_marker
forbid_loop_getter_marker = get_marker('forbid_get_event_loop')
catch_unhandled_marker = get_marker('catch_loop_exceptions')
exceptions = list()
if catch_unhandled_marker:
loop.set_exception_handler(lambda l, c: exceptions.append(c))
try:
with entrypoint(*services,
pool_size=thread_pool_size,
debug=loop_debug,
loop=loop,
**entrypoint_kwargs):
ctx = get_context(loop)
for key, value in default_context.items():
ctx[key] = value
if forbid_loop_getter_marker:
asyncio.get_event_loop.side_effect = partial(
pytest.fail, "get_event_loop is forbidden")
yield loop
if exceptions:
logging.error(
'Unhandled exceptions found:\n\n\t%s', "\n\t".join(
("Message: {m}\n\t"
"Future: {f}\n\t"
"Exception: {e}").format(m=e['message'],
f=repr(e.get('future')),
e=repr(e.get('exception')))
for e in exceptions))
pytest.fail("Unhandled exceptions found. See logs.")
finally:
with suppress(Exception):
pool.shutdown(True)
asyncio.get_event_loop.side_effect = get_event_loop
del loop
def test_aiohttp_service_sock(unix_socket_tcp):
@threaded
def http_client():
url = 'http+unix://%s/' % quote(unix_socket_tcp.getsockname(), safe='')
return requests_unixsocket.get(url).status_code
service = AIOHTTPTestApp(sock=unix_socket_tcp)
with entrypoint(service) as loop:
response = loop.run_until_complete(http_client())
assert response == 404
def test_wait_for_exception():
async def coro(arg):
await asyncio.sleep(0.1)
assert arg != 15
return arg
with entrypoint() as loop:
with pytest.raises(AssertionError):
loop.run_until_complete(wait_for(*[coro(i) for i in range(100)]))
results = loop.run_until_complete(
wait_for(*[coro(i) for i in range(17)], raise_first=False), )
assert results
assert len(results) == 17
assert isinstance(results[15], AssertionError)
assert results[:15] == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
assert results[16:] == [16]
def test_udp_server(unused_tcp_port):
class TestService(UDPServer):
DATA = []
async def handle_datagram(self, data: bytes, addr: tuple):
self.DATA.append(data)
service = TestService('127.0.0.1', unused_tcp_port)
@threaded
def writer():
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
with sock:
sock.sendto(b'hello server\n', ('127.0.0.1', unused_tcp_port))
with entrypoint(service) as loop:
loop.run_until_complete(writer())
assert TestService.DATA
assert TestService.DATA == [b'hello server\n']
def test_udp_socket_server(unix_socket_udp):
class TestService(UDPServer):
DATA = []
async def handle_datagram(self, data: bytes, addr: tuple):
self.DATA.append(data)
service = TestService(sock=unix_socket_udp)
@threaded
def writer():
sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
with sock:
sock.sendto(b'hello server\n', unix_socket_udp.getsockname())
with entrypoint(service) as loop:
loop.run_until_complete(writer())
assert TestService.DATA
assert TestService.DATA == [b'hello server\n']
def test_tcp_server(unused_tcp_port):
class TestService(TCPServer):
DATA = []
async def handle_client(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter):
self.DATA.append(await reader.readline())
writer.close()
service = TestService('127.0.0.1', unused_tcp_port)
@threaded
def writer():
with socket.create_connection(('127.0.0.1', unused_tcp_port)) as sock:
sock.send(b'hello server\n')
with entrypoint(service) as loop:
loop.run_until_complete(writer())
assert TestService.DATA
assert TestService.DATA == [b'hello server\n']
def test_tcp_server_unix(unix_socket_tcp):
class TestService(TCPServer):
DATA = []
async def handle_client(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter):
self.DATA.append(await reader.readline())
writer.close()
service = TestService(sock=unix_socket_tcp)
@threaded
def writer():
with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
sock.connect(unix_socket_tcp.getsockname())
sock.send(b'hello server\n')
with entrypoint(service) as loop:
loop.run_until_complete(writer())
assert TestService.DATA
assert TestService.DATA == [b'hello server\n']
def test_shield():
results = []
@shield
async def coro():
nonlocal results
await asyncio.sleep(0.5)
results.append(True)
async def main(loop):
task = loop.create_task(coro())
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
finally:
await asyncio.sleep(1)
with entrypoint() as loop:
loop.run_until_complete(main(loop))
assert results == [True]
def test_wait_for_cancelling():
results = []
async def coro(arg):
nonlocal results
await asyncio.sleep(0.1)
assert arg != 15
if arg > 15:
await asyncio.sleep(1)
results.append(arg)
with entrypoint() as loop:
with pytest.raises(AssertionError):
loop.run_until_complete(wait_for(*[coro(i) for i in range(100)]))
loop.run_until_complete(asyncio.sleep(2))
assert results
assert len(results) == 15
assert len(set(results)) == 15
assert frozenset(results) == frozenset(range(15))
await RobustTCPClient.start(self)
async def main(client: RPCClient):
call_count = 1000
delta = -monotonic()
for i in range(call_count):
await asyncio.gather(*[
client("multiply", x=120000, y=1000000) for _ in range(call_count)
])
delta += monotonic()
total_request_sent = call_count**2
log.info("Total executed %d requests on %.3f", total_request_sent, delta)
log.info("RPS: %.3f", total_request_sent / delta)
log.info("Close connection")
handlers = MappingProxyType({
"print": print,
})
if __name__ == "__main__":
client = RPCClient("::1", 5678, handlers=handlers)
with entrypoint(client) as loop:
loop.run_until_complete(main(client))
with io.BytesIO() as f:
response = packer.pack(result)
f.write(self.HEADER.pack(len(response)))
f.write(response)
payload = f.getvalue()
writer.write(payload)
except Exception:
writer.close()
raise
async def execute(self, method: str, kwargs: dict):
func = self.handlers[method]
if asyncio.iscoroutinefunction(func):
return await func(**kwargs)
else:
return func(**kwargs)
handlers = MappingProxyType({
'multiply': lambda x, y: x * y,
})
if __name__ == '__main__':
service = RPCServer(handlers=handlers, address='::', port=5678)
with entrypoint(service) as loop:
loop.run_forever()
async def main(host, port):
log.info('Connecting to %s:%d', host, port)
reader, writer = await asyncio.open_connection(host, port)
client = RPCClient(reader, writer)
call_count = 300
delta = -monotonic()
for i in range(call_count):
await asyncio.gather(*[
client('multiply', x=120000, y=1000000) for _ in range(call_count)
])
delta += monotonic()
total_request_sent = (call_count**2)
log.info("Total executed %d requests on %.3f", total_request_sent, delta)
log.info("RPS: %.3f", total_request_sent / delta)
await client.close()
log.info('Close connection')
if __name__ == '__main__':
with entrypoint() as loop:
loop.run_until_complete(main("::1", 5678))
p.add_argument("--host-url", required=True)
p.add_argument("--memory-tracer", action="store_true", default=False)
p.add_argument("--webhook-port", type=int, default=443)
if __name__ == "__main__":
arguments = p.parse_args()
loop = new_event_loop()
bot, manager = loop.run_until_complete(init(arguments.redis_url))
socket = bind_socket(address="0.0.0.0", port=80, proto_name="http")
services = [
TelegramWebhook(
sock=socket,
bot=bot,
manager=manager,
host=arguments.host_url,
webhooks_port=arguments.webhook_port,
),
UpdaterService(bot=bot, manager=manager),
RavenSender(sentry_dsn=sentry_key),
]
if arguments.memory_tracer:
services.append(MemoryTracer(interval=60))
with entrypoint(loop=loop, *services) as loop:
log.info("Start")
loop.run_forever()
