def test_cm_rolls_back_transaction():
conn = Mock()
tree = Tree(connection=conn)
with pytest.raises(ValueError):
with tree() as transaction: # noqa
raise ValueError()
assert conn.rollback.called
def test_cm_puts_connection_back_into_pool():
pool, conn = Mock(), Mock()
pool.getconn.return_value = conn
tree = Tree(pool=pool)
with tree() as transaction: # noqa
pass
assert pool.putconn.called
pool.putconn.assert_called_with(conn)
def run():
if config['generate_tree']:
# drop existing database and recreate
postgres_create_db(config['postgres_db'], DBNAME)
connection = psycopg2.connect(config['benchmark_db'])
tree = Tree(connection)
with tree() as transaction:
transaction.install()
# create tree with test data
generate_tree(transaction, config['levels'], config['per_level'])
connection = psycopg2.connect(config['benchmark_db'])
tree = Tree(connection)
with tree() as transaction:
postgres_analyze_db(transaction.cursor)
# build a list of benchmarks to run
benchmarks = create_benchmarks(transaction, config)
benchmarks_to_run = []
filter_benchmarks = config['filter_benchmarks']
for b in benchmarks:
if not filter_benchmarks or filter_benchmarks in b.name:
benchmarks_to_run.append(b)
print()
if len(benchmarks_to_run):
print("Running benchmarks..")
for benchmark in benchmarks_to_run:
print(benchmark.name.ljust(30), end="")
sys.stdout.flush()
duration = benchmark.run(transaction)
print(format_duration(duration))
else:
print("No benchmarks to run")
def test_close_pool():
pool = Mock()
Tree(pool=pool).close()
assert pool.closeall.called
def test_close_connection():
conn = Mock()
Tree(connection=conn).close()
assert conn.close.called
def test_it_takes_a_connection(dsn):
conn = Mock()
assert Tree(connection=conn).connection is conn
def test_cm_commits_transaction():
conn = Mock()
tree = Tree(connection=conn)
with tree() as transaction: # noqa
pass
assert conn.commit.called
def test_returned_transaction_uses_assigned_transaction_object():
conn = Mock()
tree = Tree(connection=conn)
transaction = tree.make_transaction()
assert transaction.connection is conn
def test_returned_transaction_uses_connection_from_pool():
pool, conn = Mock(), Mock()
pool.getconn.return_value = conn
tree = Tree(pool=pool)
transaction = tree.make_transaction()
assert transaction.connection is conn
def test_it_cant_deal_with_both_connection_and_pool():
with pytest.raises(TypeError):
Tree(connection=Mock(), pool=Mock())
def test_make_transaction_returns_a_transaction_object():
conn = Mock()
tree = Tree(connection=conn)
transaction = tree.make_transaction()
assert transaction.__class__ == Transaction
def test_it_requires_either_connection_or_pool():
with pytest.raises(TypeError):
Tree()
def test_it_takes_a_connection_pool(dsn):
pool = Mock()
assert Tree(pool=pool).pool is pool
def test_make_transaction_returns_a_read_write_transaction_object():
conn = Mock()
tree = Tree(connection=conn)
transaction = tree.make_transaction(write=True)
assert transaction.__class__ == ReadWriteTransaction
def test_get_connection_returns_connection_from_pool():
pool, conn = Mock(), Mock()
pool.getconn.return_value = conn
tree = Tree(pool=pool)
assert tree.get_connection() is conn
assert pool.getconn.called
def test_get_connection_returns_the_assigned_one():
conn = Mock()
tree = Tree(connection=conn)
assert tree.get_connection() is conn
