A test to use psycopg3 with anyio

test_anyio.py

"""
A quick test to try and run psycopg3 with anyio
"""

import anyio
from psycopg3 import exceptions as exc
from psycopg3.waiting import Wait, Ready
from psycopg3.connection import AsyncConnection

import logging
logger = logging.getLogger()


class AnyioAsyncConnection(AsyncConnection):
    @classmethod
    async def wait(self, gen):
        try:
            fd, s = next(gen)
            while 1:
                logger.info("fd: %s", fd)
                # TODO: what abot Wait.RW?
                if s & Wait.W:
                    await anyio.wait_socket_writable(fd)
                    fd, s = gen.send(Ready.W)
                elif s & Wait.R:
                    await anyio.wait_socket_readable(fd)
                    fd, s = gen.send(Ready.R)
                else:
                    raise exc.InternalError("bad poll status: %s")

        except StopIteration as e:
            rv = e.args[0]
            return rv


def test_connect(dsn):
    # the tests pass with PSYCOPG3_TEST_DSN="dbname=psycopg3_test host=localhost"
    # but not with PSYCOPG3_TEST_DSN="dbname=psycopg3_test"
    cnn = anyio.run(AnyioAsyncConnection.connect, dsn)
    cur = cnn.cursor()
    anyio.run(cur.execute, "select 'hello', pg_sleep(1)")
    rec = cur.fetchone()
    assert rec == ('hello', "")