psycopg2/tests/testutils.py

548 lines
16 KiB
Python
Raw Normal View History

# 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
2015-10-15 19:01:43 +03:00
import select
import operator
import platform
import unittest
from functools import wraps
from ctypes.util import find_library
2020-11-17 21:39:39 +03:00
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
2015-03-16 11:50:20 +03:00
# Silence warnings caused by the stubbornness of the Python unittest
# maintainers
2018-09-23 04:54:55 +03:00
# https://bugs.python.org/issue9424
2016-10-11 02:10:53 +03:00
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)
2018-10-23 02:39:14 +03:00
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):
2020-11-17 21:39:39 +03:00
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(
2020-11-18 00:52:11 +03:00
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.
"""
2016-12-24 02:18:22 +03:00
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:
2016-12-25 19:46:11 +03:00
self.wait(conn)
except psycopg2.OperationalError as e:
2016-12-25 19:46:11 +03:00
# 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:
2020-11-18 00:52:11 +03:00
return self.skipTest(f"replication db not configured: {e}")
2016-12-25 19:46:11 +03:00
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)
2015-10-15 19:01:43 +03:00
# 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)
2015-10-15 19:01:43 +03:00
elif state == psycopg2.extensions.POLL_WRITE:
select.select([], [pollable], [], 1)
2015-10-15 19:01:43 +03:00
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))
2016-10-11 02:10:53 +03:00
@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_
2016-10-11 02:10:53 +03:00
def skip_after_postgres(*ver):
"""Skip a test on PostgreSQL after (including) a certain version."""
ver = ver + (0,) * (3 - len(ver))
2016-10-11 02:10:53 +03:00
@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_
2016-10-11 02:10:53 +03:00
2015-06-01 19:05:11 +03:00
def libpq_version():
v = psycopg2.__libpq_version__
if v >= 90100:
v = min(v, psycopg2.extensions.libpq_version())
2015-06-01 19:05:11 +03:00
return v
2016-10-11 02:10:53 +03:00
2015-06-01 19:05:11 +03:00
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))
2016-10-11 02:10:53 +03:00
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)
2015-06-01 19:05:11 +03:00
return skip_before_libpq_
2016-10-11 02:10:53 +03:00
2015-06-01 19:05:11 +03:00
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))
2016-10-11 02:10:53 +03:00
def skip_after_libpq_(cls):
v = libpq_version()
decorator = unittest.skipIf(
v >= int("%d%02d%02d" % ver),
2020-11-18 00:52:11 +03:00
f"skipped because libpq {v}",
)
return decorator(cls)
2015-06-01 19:05:11 +03:00
return skip_after_libpq_
2016-10-11 02:10:53 +03:00
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,
2020-11-18 00:52:11 +03:00
f"skipped because Python {'.'.join(map(str, sys.version_info[:len(ver)]))}",
)
return decorator(cls)
return skip_before_python_
2016-10-11 02:10:53 +03:00
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,
2020-11-18 00:52:11 +03:00
f"skipped because Python {'.'.join(map(str, sys.version_info[:len(ver)]))}",
)
return decorator(cls)
return skip_from_python_
2016-10-11 02:10:53 +03:00
@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_
2016-10-11 02:10:53 +03:00
def skip_if_green(reason):
def skip_if_green_(cls):
decorator = unittest.skipIf(green, reason)
return decorator(cls)
return skip_if_green_
2018-10-23 02:39:14 +03:00
skip_copy_if_green = skip_if_green("copy in async mode currently not supported")
2016-10-11 02:10:53 +03:00
def skip_if_no_getrefcount(cls):
decorator = unittest.skipUnless(
hasattr(sys, 'getrefcount'),
'no sys.getrefcount()',
)
return decorator(cls)
2016-10-11 02:10:53 +03:00
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(
2020-11-18 00:52:11 +03:00
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.
"""
2020-11-17 21:39:39 +03:00
if not isinstance(reason, str):
2020-11-18 00:52:11 +03:00
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,
"batch statements": 44803,
"cancel": 41335,
"cast adds tz": 51692,
"cidr": 18846,
"composite": 27792,
"copy": 41608,
"cursor with hold": 77101,
"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,
"scroll cursor": 77102,
"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):
2020-11-17 21:39:39 +03:00
assert type is TypeError
return True
def slow(f):
2017-02-02 05:09:59 +03:00
"""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_