blob: 158fa1e3be09a86186a2503bdbfff50620401206 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
|
import pytest
import psycopg2cffi
from pgcli.main import format_output
# TODO: should this be somehow be divined from environment?
POSTGRES_USER, POSTGRES_HOST = 'postgres', 'localhost'
def db_connection(dbname=None):
conn = psycopg2cffi.connect(user=POSTGRES_USER, host=POSTGRES_HOST, dbname=dbname)
conn.autocommit = True
return conn
try:
db_connection()
CAN_CONNECT_TO_DB = True
except:
CAN_CONNECT_TO_DB = False
dbtest = pytest.mark.skipif(
not CAN_CONNECT_TO_DB,
reason="Need a postgres instance at localhost accessible by user 'postgres'")
def create_db(dbname):
with db_connection().cursor() as cur:
try:
cur.execute('''CREATE DATABASE _test_db''')
except:
pass
def drop_tables(conn):
with conn.cursor() as cur:
cur.execute('''
DROP SCHEMA public CASCADE;
CREATE SCHEMA public;
DROP SCHEMA IF EXISTS schema1 CASCADE;
DROP SCHEMA IF EXISTS schema2 CASCADE''')
def run(executor, sql, join=False):
" Return string output for the sql to be run "
result = []
for rows, headers, status in executor.run(sql):
result.extend(format_output(rows, headers, status, 'psql'))
if join:
result = '\n'.join(result)
return result
|