def test_api_expected(get_client, method, args, expected_cnts):
# Init
init_pg_world(family.dm)
client = get_client()
# Modify storage.pg.Query.upsert() to make it count it's calls per resource type
cnts = defaultdict(int)
def wrap_create_storage(create_storage):
def wrapped_create_storage(*args, **kwargs):
storage = create_storage(*args, **kwargs)
def wrap_upsert(f):
def wrapped_upsert(*args, **kwargs):
table_name = args[0]
cnts[table_name] += 1
return f(*args, **kwargs)
return wrapped_upsert
storage._q().upsert = wrap_upsert(storage._q().upsert)
return storage
return wrapped_create_storage
engine.config._config['create_storage'] = wrap_create_storage(
engine.config._config['create_storage'])
# Test
data, status, headers = getattr(client, method)(*args)
assert str(status).startswith('2')
assert dict(cnts) == expected_cnts
def test_not_null_put(get_client, status, args):
init_pg_world(cookies.dm)
client = get_client()
# from now on, we no longer accept cookies of unknown type
conn = world().storage._conn
conn.cursor().execute('ALTER TABLE cookie ALTER COLUMN type SET NOT NULL')
conn.commit()
assert client.put('cookie', 4, args)[1] == status
def _get_not_null_msg():
'''Ugly fix to make the tests work on both tested versions of PostgreSQL'''
init_pg_world(cookies.dm)
cursor = world().storage._conn.cursor()
cursor.execute('SELECT version()')
version = cursor.fetchone()[0]
if 'PostgreSQL 9' in version:
return 'null value in column "type" violates not-null constraint'
else:
return 'null value in column "type" of relation "cookie" violates not-null constraint'
def test_not_null_patch(get_client, expected_status, cookie_id, args, expected_data):
init_pg_world(cookies.dm)
client = get_client()
# from now on, we no longer accept cookies of unknown type
conn = world().storage._conn
conn.cursor().execute('ALTER TABLE cookie ALTER COLUMN type SET NOT NULL')
conn.commit()
data, status, headers = client.patch('cookie', cookie_id, args)
assert status == expected_status
print(data)
if expected_data is not None:
assert data == expected_data
def test_cookies_with_weight(get_client, status, method, args):
init_pg_world(cookies.dm)
client = get_client()
# cookies column gets additional column datamodel knows nothing about
conn = world().storage._conn
conn.cursor().execute(
'ALTER TABLE cookie ADD COLUMN weight float DEFAULT 33.33')
conn.commit()
world().storage._q()._table_columns_data = {} # delete table info cache
# get has kwargs, other have args
if method == 'get':
assert getattr(client, method)('cookie', **args)[1] == status
else:
assert getattr(client, method)('cookie', *args)[1] == status
# whatever we did, it should never influence column not in the datamodel
cur = conn.cursor()
cur.execute("SELECT DISTINCT weight FROM cookie")
data = cur.fetchall()
assert len(data) == 1
assert data[0][0] == 33.33
def init_cookies_with_user_id():
'''
Init cookies world with user_id.
Cookies 1, 2 and jar 1 are owned by user 1, others - by user 2.
'''
def set_query_class(query_cls):
'''
Make tester use CLS as query
'''
def wrap_storage(f):
def wrapped_storage(*args, **kwargs):
old_storage = f(*args, **kwargs)
new_storage = engine.PGStorage(old_storage._conn,
old_storage._schema,
query_cls=query_cls)
return new_storage
return wrapped_storage
engine.config._config['create_storage'] = wrap_storage(
engine.config._config['create_storage'])
# 1. Init world
init_pg_world(cookies.dm)
# 2. Modify cookie/jar tables
conn = engine.world().storage._conn
conn.cursor().execute('''
ALTER TABLE jar ADD COLUMN user_id integer;
ALTER TABLE cookie ADD COLUMN user_id integer;
UPDATE jar SET user_id = 2 - (id < 2)::integer;
UPDATE cookie SET user_id = 2 - (id < 3)::integer;
''')
conn.commit()
# 3. Set query class
set_query_class(UserIdQuery)
def test_disconnected_database(get_client, method, args):
init_pg_world(cookies.dm)
client = get_client()
world().storage._conn.close()
assert getattr(client, method)(*args)[1] == 500