psycopg2/tests/testutils.py

# testutils.py - utility module for psycopg2 testing.

#
# Copyright (C) 2010-2019 Daniele Varrazzo  <daniele.varrazzo@gmail.com>
# Copyright (C) 2020-2021 The Psycopg Team
#
# psycopg2 is free software: you can redistribute it and/or modify it
# under the terms of the GNU Lesser General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# In addition, as a special exception, the copyright holders give
# permission to link this program with the OpenSSL library (or with
# modified versions of OpenSSL that use the same license as OpenSSL),
# and distribute linked combinations including the two.
#
# You must obey the GNU Lesser General Public License in all respects for
# all of the code used other than OpenSSL.
#
# psycopg2 is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public
# License for more details.


import re
import os
import sys
import types
import ctypes
import select
import operator
import platform
import unittest
from functools import wraps
from ctypes.util import find_library
from io import StringIO         # noqa
from io import TextIOBase       # noqa
from importlib import reload    # noqa

import psycopg2
import psycopg2.errors
import psycopg2.extensions

from .testconfig import green, dsn, repl_dsn


# Silence warnings caused by the stubbornness of the Python unittest
# maintainers
# https://bugs.python.org/issue9424
if (not hasattr(unittest.TestCase, 'assert_')
        or unittest.TestCase.assert_ is not unittest.TestCase.assertTrue):
    # mavaff...
    unittest.TestCase.assert_ = unittest.TestCase.assertTrue
    unittest.TestCase.failUnless = unittest.TestCase.assertTrue
    unittest.TestCase.assertEquals = unittest.TestCase.assertEqual
    unittest.TestCase.failUnlessEqual = unittest.TestCase.assertEqual


def assertDsnEqual(self, dsn1, dsn2, msg=None):
    """Check that two conninfo string have the same content"""
    self.assertEqual(set(dsn1.split()), set(dsn2.split()), msg)


unittest.TestCase.assertDsnEqual = assertDsnEqual


class ConnectingTestCase(unittest.TestCase):
    """A test case providing connections for tests.

    A connection for the test is always available as `self.conn`. Others can be
    created with `self.connect()`. All are closed on tearDown.

    Subclasses needing to customize setUp and tearDown should remember to call
    the base class implementations.
    """
    def setUp(self):
        self._conns = []

    def tearDown(self):
        # close the connections used in the test
        for conn in self._conns:
            if not conn.closed:
                conn.close()

    def assertQuotedEqual(self, first, second, msg=None):
        """Compare two quoted strings disregarding eventual E'' quotes"""
        def f(s):
            if isinstance(s, str):
                return re.sub(r"\bE'", "'", s)
            elif isinstance(first, bytes):
                return re.sub(br"\bE'", b"'", s)
            else:
                return s

        return self.assertEqual(f(first), f(second), msg)

    def connect(self, **kwargs):
        try:
            self._conns
        except AttributeError as e:
            raise AttributeError(
                f"{e} (did you forget to call ConnectingTestCase.setUp()?)")

        if 'dsn' in kwargs:
            conninfo = kwargs.pop('dsn')
        else:
            conninfo = dsn
        conn = psycopg2.connect(conninfo, **kwargs)
        self._conns.append(conn)
        return conn

    def repl_connect(self, **kwargs):
        """Return a connection set up for replication

        The connection is on "PSYCOPG2_TEST_REPL_DSN" unless overridden by
        a *dsn* kwarg.

        Should raise a skip test if not available, but guard for None on
        old Python versions.
        """
        if repl_dsn is None:
            return self.skipTest("replication tests disabled by default")

        if 'dsn' not in kwargs:
            kwargs['dsn'] = repl_dsn
        try:
            conn = self.connect(**kwargs)
            if conn.async_ == 1:
                self.wait(conn)
        except psycopg2.OperationalError as e:
            # If pgcode is not set it is a genuine connection error
            # Otherwise we tried to run some bad operation in the connection
            # (e.g. bug #482) and we'd rather know that.
            if e.pgcode is None:
                return self.skipTest(f"replication db not configured: {e}")
            else:
                raise

        return conn

    def _get_conn(self):
        if not hasattr(self, '_the_conn'):
            self._the_conn = self.connect()

        return self._the_conn

    def _set_conn(self, conn):
        self._the_conn = conn

    conn = property(_get_conn, _set_conn)

    # for use with async connections only
    def wait(self, cur_or_conn):
        pollable = cur_or_conn
        if not hasattr(pollable, 'poll'):
            pollable = cur_or_conn.connection
        while True:
            state = pollable.poll()
            if state == psycopg2.extensions.POLL_OK:
                break
            elif state == psycopg2.extensions.POLL_READ:
                select.select([pollable], [], [], 1)
            elif state == psycopg2.extensions.POLL_WRITE:
                select.select([], [pollable], [], 1)
            else:
                raise Exception("Unexpected result from poll: %r", state)

    _libpq = None

    @property
    def libpq(self):
        """Return a ctypes wrapper for the libpq library"""
        if ConnectingTestCase._libpq is not None:
            return ConnectingTestCase._libpq

        libname = find_library('pq')
        if libname is None and platform.system() == 'Windows':
            raise self.skipTest("can't import libpq on windows")

        try:
            rv = ConnectingTestCase._libpq = ctypes.pydll.LoadLibrary(libname)
        except OSError as e:
            raise self.skipTest("couldn't open libpq for testing: %s" % e)
        return rv


def decorate_all_tests(obj, *decorators):
    """
    Apply all the *decorators* to all the tests defined in the TestCase *obj*.

    The decorator can also be applied to a decorator: if *obj* is a function,
    return a new decorator which can be applied either to a method or to a
    class, in which case it will decorate all the tests.
    """
    if isinstance(obj, types.FunctionType):
        def decorator(func_or_cls):
            if isinstance(func_or_cls, types.FunctionType):
                return obj(func_or_cls)
            else:
                decorate_all_tests(func_or_cls, obj)
                return func_or_cls

        return decorator

    for n in dir(obj):
        if n.startswith('test'):
            for d in decorators:
                setattr(obj, n, d(getattr(obj, n)))


@decorate_all_tests
def skip_if_no_uuid(f):
    """Decorator to skip a test if uuid is not supported by PG."""
    @wraps(f)
    def skip_if_no_uuid_(self):
        try:
            cur = self.conn.cursor()
            cur.execute("select typname from pg_type where typname = 'uuid'")
            has = cur.fetchone()
        finally:
            self.conn.rollback()

        if has:
            return f(self)
        else:
            return self.skipTest("uuid type not available on the server")

    return skip_if_no_uuid_


@decorate_all_tests
def skip_if_tpc_disabled(f):
    """Skip a test if the server has tpc support disabled."""
    @wraps(f)
    def skip_if_tpc_disabled_(self):
        cnn = self.connect()
        skip_if_crdb("2-phase commit", cnn)

        cur = cnn.cursor()
        try:
            cur.execute("SHOW max_prepared_transactions;")
        except psycopg2.ProgrammingError:
            return self.skipTest(
                "server too old: two phase transactions not supported.")
        else:
            mtp = int(cur.fetchone()[0])
        cnn.close()

        if not mtp:
            return self.skipTest(
                "server not configured for two phase transactions. "
                "set max_prepared_transactions to > 0 to run the test")
        return f(self)

    return skip_if_tpc_disabled_


def skip_before_postgres(*ver):
    """Skip a test on PostgreSQL before a certain version."""
    reason = None
    if isinstance(ver[-1], str):
        ver, reason = ver[:-1], ver[-1]

    ver = ver + (0,) * (3 - len(ver))

    @decorate_all_tests
    def skip_before_postgres_(f):
        @wraps(f)
        def skip_before_postgres__(self):
            if self.conn.info.server_version < int("%d%02d%02d" % ver):
                return self.skipTest(
                    reason or "skipped because PostgreSQL %s"
                    % self.conn.info.server_version)
            else:
                return f(self)

        return skip_before_postgres__
    return skip_before_postgres_


def skip_after_postgres(*ver):
    """Skip a test on PostgreSQL after (including) a certain version."""
    ver = ver + (0,) * (3 - len(ver))

    @decorate_all_tests
    def skip_after_postgres_(f):
        @wraps(f)
        def skip_after_postgres__(self):
            if self.conn.info.server_version >= int("%d%02d%02d" % ver):
                return self.skipTest("skipped because PostgreSQL %s"
                    % self.conn.info.server_version)
            else:
                return f(self)

        return skip_after_postgres__
    return skip_after_postgres_


def libpq_version():
    v = psycopg2.__libpq_version__
    if v >= 90100:
        v = min(v, psycopg2.extensions.libpq_version())
    return v


def skip_before_libpq(*ver):
    """Skip a test if libpq we're linked to is older than a certain version."""
    ver = ver + (0,) * (3 - len(ver))

    def skip_before_libpq_(cls):
        v = libpq_version()
        decorator = unittest.skipIf(
            v < int("%d%02d%02d" % ver),
            f"skipped because libpq {v}",
        )
        return decorator(cls)
    return skip_before_libpq_


def skip_after_libpq(*ver):
    """Skip a test if libpq we're linked to is newer than a certain version."""
    ver = ver + (0,) * (3 - len(ver))

    def skip_after_libpq_(cls):
        v = libpq_version()
        decorator = unittest.skipIf(
            v >= int("%d%02d%02d" % ver),
            f"skipped because libpq {v}",
        )
        return decorator(cls)
    return skip_after_libpq_


def skip_before_python(*ver):
    """Skip a test on Python before a certain version."""
    def skip_before_python_(cls):
        decorator = unittest.skipIf(
            sys.version_info[:len(ver)] < ver,
            f"skipped because Python {'.'.join(map(str, sys.version_info[:len(ver)]))}",
        )
        return decorator(cls)
    return skip_before_python_


def skip_from_python(*ver):
    """Skip a test on Python after (including) a certain version."""
    def skip_from_python_(cls):
        decorator = unittest.skipIf(
            sys.version_info[:len(ver)] >= ver,
            f"skipped because Python {'.'.join(map(str, sys.version_info[:len(ver)]))}",
        )
        return decorator(cls)
    return skip_from_python_


@decorate_all_tests
def skip_if_no_superuser(f):
    """Skip a test if the database user running the test is not a superuser"""
    @wraps(f)
    def skip_if_no_superuser_(self):
        try:
            return f(self)
        except psycopg2.errors.InsufficientPrivilege:
            self.skipTest("skipped because not superuser")

    return skip_if_no_superuser_


def skip_if_green(reason):
    def skip_if_green_(cls):
        decorator = unittest.skipIf(green, reason)
        return decorator(cls)
    return skip_if_green_


skip_copy_if_green = skip_if_green("copy in async mode currently not supported")


def skip_if_no_getrefcount(cls):
    decorator = unittest.skipUnless(
        hasattr(sys, 'getrefcount'),
        'no sys.getrefcount()',
    )
    return decorator(cls)


def skip_if_windows(cls):
    """Skip a test if run on windows"""
    decorator = unittest.skipIf(
        platform.system() == 'Windows',
        "Not supported on Windows",
    )
    return decorator(cls)


def crdb_version(conn, __crdb_version=[]):
    """
    Return the CockroachDB version if that's the db being tested, else None.

    Return the number as an integer similar to PQserverVersion: return
    v20.1.3 as 200103.

    Assume all the connections are on the same db: return a cached result on
    following calls.

    """
    if __crdb_version:
        return __crdb_version[0]

    sver = conn.info.parameter_status("crdb_version")
    if sver is None:
        __crdb_version.append(None)
    else:
        m = re.search(r"\bv(\d+)\.(\d+)\.(\d+)", sver)
        if not m:
            raise ValueError(
                f"can't parse CockroachDB version from {sver}")

        ver = int(m.group(1)) * 10000 + int(m.group(2)) * 100 + int(m.group(3))
        __crdb_version.append(ver)

    return __crdb_version[0]


def skip_if_crdb(reason, conn=None, version=None):
    """Skip a test or test class if we are testing against CockroachDB.

    Can be used as a decorator for tests function or classes:

        @skip_if_crdb("my reason")
        class SomeUnitTest(UnitTest):
            # ...

    Or as a normal function if the *conn* argument is passed.

    If *version* is specified it should be a string such as ">= 20.1", "< 20",
    "== 20.1.3": the test will be skipped only if the version matches.

    """
    if not isinstance(reason, str):
        raise TypeError(f"reason should be a string, got {reason!r} instead")

    if conn is not None:
        ver = crdb_version(conn)
        if ver is not None and _crdb_match_version(ver, version):
            if reason in crdb_reasons:
                reason = (
                    "%s (https://github.com/cockroachdb/cockroach/issues/%s)"
                    % (reason, crdb_reasons[reason]))
            raise unittest.SkipTest(
                f"not supported on CockroachDB {ver}: {reason}")

    @decorate_all_tests
    def skip_if_crdb_(f):
        @wraps(f)
        def skip_if_crdb__(self, *args, **kwargs):
            skip_if_crdb(reason, conn=self.connect(), version=version)
            return f(self, *args, **kwargs)

        return skip_if_crdb__

    return skip_if_crdb_


# mapping from reason description to ticket number
crdb_reasons = {
    "2-phase commit": 22329,
    "backend pid": 35897,
    "cancel": 41335,
    "cast adds tz": 51692,
    "cidr": 18846,
    "composite": 27792,
    "copy": 41608,
    "deferrable": 48307,
    "encoding": 35882,
    "hstore": 41284,
    "infinity date": 41564,
    "interval style": 35807,
    "large objects": 243,
    "named cursor": 41412,
    "nested array": 32552,
    "notify": 41522,
    "password_encryption": 42519,
    "range": 41282,
    "stored procedure": 1751,
}


def _crdb_match_version(version, pattern):
    if pattern is None:
        return True

    m = re.match(r'^(>|>=|<|<=|==|!=)\s*(\d+)(?:\.(\d+))?(?:\.(\d+))?$', pattern)
    if m is None:
        raise ValueError(
            "bad crdb version pattern %r: should be 'OP MAJOR[.MINOR[.BUGFIX]]'"
            % pattern)

    ops = {'>': 'gt', '>=': 'ge', '<': 'lt', '<=': 'le', '==': 'eq', '!=': 'ne'}
    op = getattr(operator, ops[m.group(1)])
    ref = int(m.group(2)) * 10000 + int(m.group(3) or 0) * 100 + int(m.group(4) or 0)
    return op(version, ref)


class raises_typeerror:
    def __enter__(self):
        pass

    def __exit__(self, type, exc, tb):
        assert type is TypeError
        return True


def slow(f):
    """Decorator to mark slow tests we may want to skip

    Note: in order to find slow tests you can run:

    make check 2>&1 | ts -i "%.s" | sort -n
    """
    @wraps(f)
    def slow_(self):
        if os.environ.get('PSYCOPG2_TEST_FAST', '0') != '0':
            return self.skipTest("slow test")
        return f(self)
    return slow_


def restore_types(f):
    """Decorator to restore the adaptation system after running a test"""
    @wraps(f)
    def restore_types_(self):
        types = psycopg2.extensions.string_types.copy()
        adapters = psycopg2.extensions.adapters.copy()
        try:
            return f(self)
        finally:
            psycopg2.extensions.string_types.clear()
            psycopg2.extensions.string_types.update(types)
            psycopg2.extensions.adapters.clear()
            psycopg2.extensions.adapters.update(adapters)

    return restore_types_