How to use the asyncpg._testbase.ConnectedTestCase function in asyncpg

# Copyright (C) 2016-present the asyncpg authors and contributors
# 
#
# This module is part of asyncpg and is released under
# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0


import asyncpg

from asyncpg import _testbase as tb


class TestTransaction(tb.ConnectedTestCase):

    async def test_transaction_regular(self):
        self.assertIsNone(self.con._top_xact)
        self.assertFalse(self.con.is_in_transaction())
        tr = self.con.transaction()
        self.assertIsNone(self.con._top_xact)
        self.assertFalse(self.con.is_in_transaction())

        with self.assertRaises(ZeroDivisionError):
            async with tr as with_tr:
                self.assertIs(self.con._top_xact, tr)
                self.assertTrue(self.con.is_in_transaction())

                # We don't return the transaction object from __aenter__,
                # to make it harder for people to use '.rollback()' and
                # '.commit()' from within an 'async with' block.

await con1.remove_listener('12+"34', listener1)

    async def test_dangling_listener_warns(self):
        async with self.create_pool(database='postgres') as pool:
            with self.assertWarnsRegex(
                    exceptions.InterfaceWarning,
                    '.*Connection.*is being released to the pool but '
                    'has 1 active notification listener'):
                async with pool.acquire() as con:
                    def listener1(*args):
                        pass

                    await con.add_listener('ipc', listener1)


class TestLogListeners(tb.ConnectedTestCase):

    @tb.with_connection_options(server_settings={
        'client_min_messages': 'notice'
    })
    async def test_log_listener_01(self):
        q1 = asyncio.Queue()

        def notice_callb(con, message):
            # Message fields depend on PG version, hide some values.
            dct = message.as_dict()
            del dct['server_source_line']
            q1.put_nowait((con, type(message), dct))

        async def raise_notice():
            await self.con.execute(
                """DO $$

'must be greater than zero'):
                    async for _ in st.cursor(prefetch=prefetch):  # NOQA
                        pass

    async def test_cursor_iterable_06(self):
        recs = []

        async with self.con.transaction():
            async for rec in self.con.cursor(
                    'SELECT generate_series(0, $1::int)', 10):
                recs.append(rec)

        self.assertEqual(recs, [(i,) for i in range(11)])


class TestCursor(tb.ConnectedTestCase):

    async def test_cursor_01(self):
        st = await self.con.prepare('SELECT generate_series(0, 20)')
        with self.assertRaisesRegex(asyncpg.NoActiveSQLTransactionError,
                                    'cursor cannot be created.*transaction'):
            await st.cursor()

    async def test_cursor_02(self):
        st = await self.con.prepare('SELECT generate_series(0, 20)')
        async with self.con.transaction():
            cur = await st.cursor()

            for i in range(-1, 1):
                with self.assertRaisesRegex(asyncpg.InterfaceError,
                                            'greater than zero'):
                    await cur.fetch(i)

os.chmod(passdir, stat.S_IRWXU)

    async def test_connect_args_validation(self):
        for val in {-1, 'a', True, False, 0}:
            with self.assertRaisesRegex(ValueError, 'greater than 0'):
                await asyncpg.connect(command_timeout=val)

        for arg in {'max_cacheable_statement_size',
                    'max_cached_statement_lifetime',
                    'statement_cache_size'}:
            for val in {None, -1, True, False}:
                with self.assertRaisesRegex(ValueError, 'greater or equal'):
                    await asyncpg.connect(**{arg: val})


class TestConnection(tb.ConnectedTestCase):

    async def test_connection_isinstance(self):
        self.assertTrue(isinstance(self.con, connection.Connection))
        self.assertTrue(isinstance(self.con, object))
        self.assertFalse(isinstance(self.con, list))

    async def test_connection_use_after_close(self):
        def check():
            return self.assertRaisesRegex(asyncpg.InterfaceError,
                                          'connection is closed')

        await self.con.close()

        with check():
            await self.con.add_listener('aaa', lambda: None)

# Copyright (C) 2016-present the asyncpg authors and contributors
# 
#
# This module is part of asyncpg and is released under
# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0


import asyncpg
from asyncpg import _testbase as tb


class TestExceptions(tb.ConnectedTestCase):

    def test_exceptions_exported(self):
        for err in ('PostgresError', 'SubstringError', 'InterfaceError'):
            self.assertTrue(hasattr(asyncpg, err))
            self.assertIn(err, asyncpg.__all__)

        for err in ('PostgresMessage',):
            self.assertFalse(hasattr(asyncpg, err))
            self.assertNotIn(err, asyncpg.__all__)

        self.assertIsNone(asyncpg.PostgresError.schema_name)

    async def test_exceptions_unpacking(self):
        try:
            await self.con.execute('SELECT * FROM _nonexistent_')
        except asyncpg.UndefinedTableError as e:

#
# This module is part of asyncpg and is released under
# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0


import asyncio
import datetime
import io
import tempfile

import asyncpg
from asyncpg import _testbase as tb
from asyncpg import compat


class TestCopyFrom(tb.ConnectedTestCase):

    async def test_copy_from_table_basics(self):
        await self.con.execute('''
            CREATE TABLE copytab(a text, "b~" text, i int);
            INSERT INTO copytab (a, "b~", i) (
                SELECT 'a' || i::text, 'b' || i::text, i
                FROM generate_series(1, 5) AS i
            );
            INSERT INTO copytab (a, "b~", i) VALUES('*', NULL, NULL);
        ''')

        try:
            f = io.BytesIO()

            # Basic functionality.
            res = await self.con.copy_from_table('copytab', output=f)

# Copyright (C) 2016-present the asyncpg authors and contributors
# 
#
# This module is part of asyncpg and is released under
# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0


import asyncpg
from asyncpg import _testbase as tb

ERRNUM = 'unexpected number of attributes of composite type'
ERRTYP = 'unexpected data type of composite type'


class TestCacheInvalidation(tb.ConnectedTestCase):

    def _get_cached_statements(self, connection=None):
        if connection is None:
            connection = self.con
        return list(connection._stmt_cache.iter_statements())

    def _check_statements_are_not_closed(self, statements):
        self.assertGreater(len(statements), 0)
        self.assertTrue(all(not s.closed for s in statements))

    def _check_statements_are_closed(self, statements):
        self.assertGreater(len(statements), 0)
        self.assertTrue(all(s.closed for s in statements))

    async def test_prepare_cache_invalidation_silent(self):
        await self.con.execute('CREATE TABLE tab1(a int, b int)')

for methname in {'fetch', 'fetchrow', 'fetchval', 'execute'}:
            with self.assertRaises(asyncio.TimeoutError), \
                    self.assertRunUnder(MAX_RUNTIME):
                meth = getattr(self.con, methname)
                await meth('select pg_sleep(10)')
            self.assertEqual(await self.con.fetch('select 1'), [(1,)])


class SlowPrepareConnection(pg_connection.Connection):
    """Connection class to test timeouts."""
    async def _get_statement(self, query, timeout):
        await asyncio.sleep(0.3)
        return await super()._get_statement(query, timeout)


class TestTimeoutCoversPrepare(tb.ConnectedTestCase):

    @tb.with_connection_options(connection_class=SlowPrepareConnection,
                                command_timeout=0.3)
    async def test_timeout_covers_prepare_01(self):
        for methname in {'fetch', 'fetchrow', 'fetchval', 'execute'}:
            with self.assertRaises(asyncio.TimeoutError):
                meth = getattr(self.con, methname)
                await meth('select pg_sleep($1)', 0.2)

# Copyright (C) 2016-present the asyncpg authors and contributors
# 
#
# This module is part of asyncpg and is released under
# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0


import asyncio
import asyncpg

from asyncpg import _testbase as tb


class TestExecuteScript(tb.ConnectedTestCase):

    async def test_execute_script_1(self):
        self.assertEqual(self.con._protocol.queries_count, 0)
        status = await self.con.execute('''
            SELECT 1;

            SELECT true FROM pg_type WHERE false = true;

            SELECT generate_series(0, 9);
        ''')
        self.assertEqual(self.con._protocol.queries_count, 1)
        self.assertEqual(status, 'SELECT 10')

    async def test_execute_script_2(self):
        status = await self.con.execute('''
            CREATE TABLE mytab (a int);

#
# This module is part of asyncpg and is released under
# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0


import asyncio

import asyncpg
from asyncpg import connection as pg_connection
from asyncpg import _testbase as tb


MAX_RUNTIME = 0.5


class TestTimeout(tb.ConnectedTestCase):

    async def test_timeout_01(self):
        for methname in {'fetch', 'fetchrow', 'fetchval', 'execute'}:
            with self.assertRaises(asyncio.TimeoutError), \
                    self.assertRunUnder(MAX_RUNTIME):
                meth = getattr(self.con, methname)
                await meth('select pg_sleep(10)', timeout=0.02)
            self.assertEqual(await self.con.fetch('select 1'), [(1,)])

    async def test_timeout_02(self):
        st = await self.con.prepare('select pg_sleep(10)')

        for methname in {'fetch', 'fetchrow', 'fetchval'}:
            with self.assertRaises(asyncio.TimeoutError), \
                    self.assertRunUnder(MAX_RUNTIME):
                meth = getattr(st, methname)