108 lines
3.2 KiB
Python
108 lines
3.2 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
import psycopg2
|
|
|
|
|
|
class StubPsycopg2(list):
|
|
"""A stubbed version of the psycopg2 module, used to implement PostgreSQL
|
|
server interaction tests, without having to talk to a running
|
|
PostgreSQL server instance."""
|
|
def __init__(self, *connections):
|
|
self.conn_params = None
|
|
super().__init__(connections)
|
|
|
|
def add(self, connection):
|
|
self.append(connection)
|
|
return self
|
|
|
|
def add_auth_failure(self):
|
|
return self.add(psycopg2.OperationalError(
|
|
'FATAL: password authentication failed for user "test"'))
|
|
|
|
def add_conn_timeout(self):
|
|
return self.add(psycopg2.OperationalError('timeout expired'))
|
|
|
|
def add_conn_failure(self):
|
|
return self.add(psycopg2.OperationalError(
|
|
'could not connect to server: No route to host / Is the ' +
|
|
'server running on host "..." and accepting / TCP/IP ' +
|
|
'connections on port ...?'))
|
|
return self
|
|
|
|
def add_admin_startup(self):
|
|
return self.add(psycopg2.errors.OperationalError(
|
|
'FATAL: the database system is starting up'))
|
|
|
|
def add_admin_shutdown(self):
|
|
return self.add(psycopg2.errors.AdminShutdown(
|
|
'terminating connection due to administrator command'))
|
|
|
|
def add_connection(self, connection):
|
|
return self.add(connection)
|
|
|
|
def connect(self, **conn_params):
|
|
self.conn_params = conn_params
|
|
response = self.pop(0)
|
|
if isinstance(response, Exception):
|
|
raise response
|
|
return response
|
|
|
|
|
|
class StubConnection(list):
|
|
"""A stubbed version of a psycopg2 connection, used to implement
|
|
PostgreSQL server interaction tests, without having to talk to a
|
|
running PostgreSQL server instance."""
|
|
def __init__(self, *cursors):
|
|
super().__init__(cursors)
|
|
self.connected = True
|
|
self.autocommit = False
|
|
|
|
def close(self):
|
|
self.connected = False
|
|
|
|
def cursor(self):
|
|
if not self.connected:
|
|
raise psycopg2.InterfaceError("connection already closed")
|
|
response = self.pop(0)
|
|
if isinstance(response, Exception):
|
|
raise response
|
|
return response
|
|
|
|
|
|
class StubCursor(list):
|
|
"""A stubbed version of a psycopg2 cursor, used to implement PostgreSQL
|
|
server interaction tests, without having to talk to a running
|
|
PostgreSQL server instance."""
|
|
def __init__(self, *results):
|
|
super().__init__(results)
|
|
self.statusmessage = None
|
|
self.rows = None
|
|
self.query = None
|
|
self.params = None
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, a, b, c):
|
|
pass
|
|
|
|
def execute(self, query, params=()):
|
|
self.query = query
|
|
self.params = params
|
|
response = self.pop(0)
|
|
if isinstance(response, Exception):
|
|
raise response
|
|
if callable(response):
|
|
return response(query)
|
|
status, rows = response
|
|
self.statusmessage = status
|
|
self.rows = rows if rows is not None else []
|
|
|
|
def fetchone(self):
|
|
if self.rows is None:
|
|
raise psycopg2.ProgrammingError("no results to fetch")
|
|
try:
|
|
return self.rows.pop(0)
|
|
except IndexError:
|
|
return None
|