# coding: utf-8
from __future__ import unicode_literals
from collections import OrderedDict
import os
import unittest
from datetime import datetime
from sqlalchemy import FLOAT, TEXT, BIGINT
from sqlalchemy.exc import IntegrityError, SQLAlchemyError, ArgumentError
from dataset import connect, chunked
from .sample_data import TEST_DATA, TEST_CITY_1
class DatabaseTestCase(unittest.TestCase):
def setUp(self):
os.environ.setdefault('DATABASE_URL', 'sqlite:///:memory:')
self.db = connect(os.environ['DATABASE_URL'])
self.tbl = self.db['weather']
self.tbl.insert_many(TEST_DATA)
def tearDown(self):
for table in self.db.tables:
self.db[table].drop()
def test_valid_database_url(self):
assert self.db.url, os.environ['DATABASE_URL']
def test_database_url_query_string(self):
db = connect('sqlite:///:memory:/?cached_statements=1')
assert 'cached_statements' in db.url, db.url
def test_tables(self):
assert self.db.tables == ['weather'], self.db.tables
def test_contains(self):
assert 'weather' in self.db, self.db.tables
def test_create_table(self):
table = self.db['foo']
assert table.table.exists()
assert len(table.table.columns) == 1, table.table.columns
assert 'id' in table.table.c, table.table.c
def test_create_table_no_ids(self):
if 'mysql' in self.db.engine.dialect.dbapi.__name__:
return
if 'sqlite' in self.db.engine.dialect.dbapi.__name__:
return
table = self.db.create_table("foo_no_id", primary_id=False)
assert table.table.exists()
assert len(table.table.columns) == 0, table.table.columns
def test_create_table_custom_id1(self):
pid = "string_id"
table = self.db.create_table("foo2", pid, self.db.types.string(255))
assert table.table.exists()
assert len(table.table.columns) == 1, table.table.columns
assert pid in table.table.c, table.table.c
table.insert({pid: 'foobar'})
assert table.find_one(string_id='foobar')[pid] == 'foobar'
def test_create_table_custom_id2(self):
pid = "string_id"
table = self.db.create_table("foo3", pid, self.db.types.string(50))
assert table.table.exists()
assert len(table.table.columns) == 1, table.table.columns
assert pid in table.table.c, table.table.c
table.insert({pid: 'foobar'})
assert table.find_one(string_id='foobar')[pid] == 'foobar'
def test_create_table_custom_id3(self):
pid = "int_id"
table = self.db.create_table("foo4", primary_id=pid)
assert table.table.exists()
assert len(table.table.columns) == 1, table.table.columns
assert pid in table.table.c, table.table.c
table.insert({pid: 123})
table.insert({pid: 124})
assert table.find_one(int_id=123)[pid] == 123
assert table.find_one(int_id=124)[pid] == 124
self.assertRaises(IntegrityError, lambda: table.insert({pid: 123}))
def test_create_table_shorthand1(self):
pid = "int_id"
table = self.db.get_table('foo5', pid)
assert table.table.exists
assert len(table.table.columns) == 1, table.table.columns
assert pid in table.table.c, table.table.c
table.insert({'int_id': 123})
table.insert({'int_id': 124})
assert table.find_one(int_id=123)['int_id'] == 123
assert table.find_one(int_id=124)['int_id'] == 124
self.assertRaises(IntegrityError, lambda: table.insert({'int_id': 123}))
def test_create_table_shorthand2(self):
pid = "string_id"
table = self.db.get_table('foo6', primary_id=pid,
primary_type=self.db.types.string(255))
assert table.table.exists
assert len(table.table.columns) == 1, table.table.columns
assert pid in table.table.c, table.table.c
table.insert({
'string_id': 'foobar'})
assert table.find_one(string_id='foobar')['string_id'] == 'foobar'
def test_with(self):
init_length = len(self.db['weather'])
with self.assertRaises(ValueError):
with self.db as tx:
tx['weather'].insert({'date': datetime(2011, 1, 1),
'temperature': 1,
'place': 'tmp_place'})
raise ValueError()
assert len(self.db['weather']) == init_length
def test_invalid_values(self):
if 'mysql' in self.db.engine.dialect.dbapi.__name__:
# WARNING: mysql seems to be doing some weird type casting upon insert.
# The mysql-python driver is not affected but it isn't compatible with Python 3
# Conclusion: use postgresql.
return
with self.assertRaises(SQLAlchemyError):
tbl = self.db['weather']
tbl.insert({'date': True, 'temperature': 'wrong_value', 'place': 'tmp_place'})
def test_load_table(self):
tbl = self.db.load_table('weather')
assert tbl.table.name == self.tbl.table.name
def test_query(self):
r = self.db.query('SELECT COUNT(*) AS num FROM weather').next()
assert r['num'] == len(TEST_DATA), r
def test_table_cache_updates(self):
tbl1 = self.db.get_table('people')
data = OrderedDict([('first_name', 'John'), ('last_name', 'Smith')])
tbl1.insert(data)
data['id'] = 1
tbl2 = self.db.get_table('people')
assert dict(tbl2.all().next()) == dict(data), (tbl2.all().next(), data)
class TableTestCase(unittest.TestCase):
def setUp(self):
self.db = connect('sqlite:///:memory:')
self.tbl = self.db['weather']
for row in TEST_DATA:
self.tbl.insert(row)
def test_insert(self):
assert len(self.tbl) == len(TEST_DATA), len(self.tbl)
last_id = self.tbl.insert({
'date': datetime(2011, 1, 2),
'temperature': -10,
'place': 'Berlin'}
)
assert len(self.tbl) == len(TEST_DATA) + 1, len(self.tbl)
assert self.tbl.find_one(id=last_id)['place'] == 'Berlin'
def test_insert_ignore(self):
self.tbl.insert_ignore({
'date': datetime(2011, 1, 2),
'temperature': -10,
'place': 'Berlin'},
['place']
)
assert len(self.tbl) == len(TEST_DATA) + 1, len(self.tbl)
self.tbl.insert_ignore({
'date': datetime(2011, 1, 2),
'temperature': -10,
'place': 'Berlin'},
['place']
)
assert len(self.tbl) == len(TEST_DATA) + 1, len(self.tbl)
def test_insert_ignore_all_key(self):
for i in range(0, 4):
self.tbl.insert_ignore({
'date': datetime(2011, 1, 2),
'temperature': -10,
'place': 'Berlin'},
['date', 'temperature', 'place']
)
assert len(self.tbl) == len(TEST_DATA) + 1, len(self.tbl)
def test_upsert(self):
self.tbl.upsert({
'date': datetime(2011, 1, 2),
'temperature': -10,
'place': 'Berlin'},
['place']
)
assert len(self.tbl) == len(TEST_DATA) + 1, len(self.tbl)
self.tbl.upsert({
'date': datetime(2011, 1, 2),
'temperature': -10,
'place': 'Berlin'},
['place']
)
assert len(self.tbl) == len(TEST_DATA) + 1, len(self.tbl)
def test_upsert_single_column(self):
table = self.db['banana_single_col']
table.upsert({
'color': 'Yellow'},
['color']
)
assert len(table) == 1, len(table)
table.upsert({
'color': 'Yellow'},
['color']
)
assert len(table) == 1, len(table)
def test_upsert_all_key(self):
assert len(self.tbl) == len(TEST_DATA), len(self.tbl)
for i in range(0, 2):
self.tbl.upsert({
'date': datetime(2011, 1, 2),
'temperature': -10,
'place': 'Berlin'},
['date', 'temperature', 'place']
)
assert len(self.tbl) == len(TEST_DATA) + 1, len(self.tbl)
def test_upsert_id(self):
table = self.db['banana_with_id']
data = dict(id=10, title='I am a banana!')
table.upsert(data, ['id'])
assert len(table) == 1, len(table)
def test_update_while_iter(self):
for row in self.tbl:
row['foo'] = 'bar'
self.tbl.update(row, ['place', 'date'])
assert len(self.tbl) == len(TEST_DATA), len(self.tbl)
def test_weird_column_names(self):
with self.assertRaises(ValueError):
self.tbl.insert({
'date': datetime(2011, 1, 2),
'temperature': -10,
'foo.bar': 'Berlin',
'qux.bar': 'Huhu'
})
def test_cased_column_names(self):
tbl = self.db['cased_column_names']
tbl.insert({
'place': 'Berlin',
})
tbl.insert({
'Place': 'Berlin',
})
tbl.insert({
'PLACE ': 'Berlin',
})
assert len(tbl.columns) == 2, tbl.columns
assert len(list(tbl.find(Place='Berlin'))) == 3
assert len(list(tbl.find(place='Berlin'))) == 3
assert len(list(tbl.find(PLACE='Berlin'))) == 3
def test_invalid_column_names(self):
tbl = self.db['weather']
with self.assertRaises(ValueError):
tbl.insert({None: 'banana'})
with self.assertRaises(ValueError):
tbl.insert({'': 'banana'})
with self.assertRaises(ValueError):
tbl.insert({'-': 'banana'})
def test_delete(self):
self.tbl.insert({
'date': datetime(2011, 1, 2),
'temperature': -10,
'place': 'Berlin'}
)
original_count = len(self.tbl)
assert len(self.tbl) == len(TEST_DATA) + 1, len(self.tbl)
# Test bad use of API
with self.assertRaises(ArgumentError):
self.tbl.delete({'place': 'Berlin'})
assert len(self.tbl) == original_count, len(self.tbl)
assert self.tbl.delete(place='Berlin') is True, 'should return 1'
assert len(self.tbl) == len(TEST_DATA), len(self.tbl)
assert self.tbl.delete() is True, 'should return non zero'
assert len(self.tbl) == 0, len(self.tbl)
def test_repr(self):
assert repr(self.tbl) == '<Table(weather)>', \
'the representation should be <Table(weather)>'
def test_delete_nonexist_entry(self):
assert self.tbl.delete(place='Berlin') is False, \
'entry not exist, should fail to delete'
def test_find_one(self):
self.tbl.insert({
'date': datetime(2011, 1, 2),
'temperature': -10,
'place': 'Berlin'
})
d = self.tbl.find_one(place='Berlin')
assert d['temperature'] == -10, d
d = self.tbl.find_one(place='Atlantis')
assert d is None, d
def test_count(self):
assert len(self.tbl) == 6, len(self.tbl)
length = self.tbl.count(place=TEST_CITY_1)
assert length == 3, length
def test_find(self):
ds = list(self.tbl.find(place=TEST_CITY_1))
assert len(ds) == 3, ds
ds = list(self.tbl.find(place=TEST_CITY_1, _limit=2))
assert len(ds) == 2, ds
ds = list(self.tbl.find(place=TEST_CITY_1, _limit=2, _step=1))
assert len(ds) == 2, ds
ds = list(self.tbl.find(place=TEST_CITY_1, _limit=1, _step=2))
assert len(ds) == 1, ds
ds = list(self.tbl.find(_step=2))
assert len(ds) == len(TEST_DATA), ds
ds = list(self.tbl.find(order_by=['temperature']))
assert ds[0]['temperature'] == -1, ds
ds = list(self.tbl.find(order_by=['-temperature']))
assert ds[0]['temperature'] == 8, ds
ds = list(self.tbl.find(self.tbl.table.columns.temperature > 4))
assert len(ds) == 3, ds
def test_find_dsl(self):
ds = list(self.tbl.find(place={'like': '%lw%'}))
assert len(ds) == 3, ds
ds = list(self.tbl.find(temperature={'>': 5}))
assert len(ds) == 2, ds
ds = list(self.tbl.find(temperature={'>=': 5}))
assert len(ds) == 3, ds
ds = list(self.tbl.find(temperature={'<': 0}))
assert len(ds) == 1, ds
ds = list(self.tbl.find(temperature={'<=': 0}))
assert len(ds) == 2, ds
ds = list(self.tbl.find(temperature={'!=': -1}))
assert len(ds) == 5, ds
ds = list(self.tbl.find(temperature={'between': [5, 8]}))
assert len(ds) == 3, ds
ds = list(self.tbl.find(place={'=': 'G€lway'}))
assert len(ds) == 3, ds
ds = list(self.tbl.find(place={'ilike': '%LwAy'}))
assert len(ds) == 3, ds
def test_offset(self):
ds = list(self.tbl.find(place=TEST_CITY_1, _offset=1))
assert len(ds) == 2, ds
ds = list(self.tbl.find(place=TEST_CITY_1, _limit=2, _offset=2))
assert len(ds) == 1, ds
def test_streamed(self):
ds = list(self.tbl.find(place=TEST_CITY_1, _streamed=True, _step=1))
assert len(ds) == 3, len(ds)
for row in self.tbl.find(place=TEST_CITY_1, _streamed=True, _step=1):
row['temperature'] = -1
self.tbl.update(row, ['id'])
def test_distinct(self):
x = list(self.tbl.distinct('place'))
assert len(x) == 2, x
x = list(self.tbl.distinct('place', 'date'))
assert len(x) == 6, x
x = list(self.tbl.distinct(
'place', 'date',
self.tbl.table.columns.date >= datetime(2011, 1, 2, 0, 0)))
assert len(x) == 4, x
x = list(self.tbl.distinct('temperature', place='B€rkeley'))
assert len(x) == 3, x
x = list(self.tbl.distinct('temperature',
place=['B€rkeley', 'G€lway']))
assert len(x) == 6, x
def test_insert_many(self):
data = TEST_DATA * 100
self.tbl.insert_many(data, chunk_size=13)
assert len(self.tbl) == len(data) + 6
def test_chunked_insert(self):
data = TEST_DATA * 100
with chunked.ChunkedInsert(self.tbl) as chunk_tbl:
for item in data:
chunk_tbl.insert(item)
assert len(self.tbl) == len(data) + 6
def test_chunked_insert_callback(self):
data = TEST_DATA * 100
N = 0
def callback(queue):
nonlocal N
N += len(queue)
with chunked.ChunkedInsert(self.tbl, callback=callback) as chunk_tbl:
for item in data:
chunk_tbl.insert(item)
assert len(data) == N
assert len(self.tbl) == len(data) + 6
def test_update_many(self):
tbl = self.db['update_many_test']
tbl.insert_many([
dict(temp=10), dict(temp=20), dict(temp=30)
])
tbl.update_many(
[dict(id=1, temp=50), dict(id=3, temp=50)], 'id'
)
# Ensure data has been updated.
assert tbl.find_one(id=1)['temp'] == tbl.find_one(id=3)['temp']
def test_upsert_many(self):
# Also tests updating on records with different attributes
tbl = self.db['upsert_many_test']
W = 100
tbl.upsert_many([dict(age=10), dict(weight=W)], 'id')
assert tbl.find_one(id=1)['age'] == 10
tbl.upsert_many([dict(id=1, age=70), dict(id=2, weight=W / 2)], 'id')
assert tbl.find_one(id=2)['weight'] == W / 2
def test_drop_operations(self):
assert self.tbl._table is not None, \
'table shouldn\'t be dropped yet'
self.tbl.drop()
assert self.tbl._table is None, \
'table should be dropped now'
assert list(self.tbl.all()) == [], self.tbl.all()
assert self.tbl.count() == 0, self.tbl.count()
def test_table_drop(self):
assert 'weather' in self.db
self.db['weather'].drop()
assert 'weather' not in self.db
def test_table_drop_then_create(self):
assert 'weather' in self.db
self.db['weather'].drop()
assert 'weather' not in self.db
self.db['weather'].insert({'foo': 'bar'})
def test_columns(self):
cols = self.tbl.columns
assert len(list(cols)) == 4, 'column count mismatch'
assert 'date' in cols and 'temperature' in cols and 'place' in cols
def test_drop_column(self):
try:
self.tbl.drop_column('date')
assert 'date' not in self.tbl.columns
except RuntimeError:
pass
def test_iter(self):
c = 0
for row in self.tbl:
c += 1
assert c == len(self.tbl)
def test_update(self):
date = datetime(2011, 1, 2)
res = self.tbl.update({
'date': date,
'temperature': -10,
'place': TEST_CITY_1},
['place', 'date']
)
assert res, 'update should return True'
m = self.tbl.find_one(place=TEST_CITY_1, date=date)
assert m['temperature'] == -10, \
'new temp. should be -10 but is %d' % m['temperature']
def test_create_column(self):
tbl = self.tbl
tbl.create_column('foo', FLOAT)
assert 'foo' in tbl.table.c, tbl.table.c
assert isinstance(tbl.table.c['foo'].type, FLOAT), \
tbl.table.c['foo'].type
assert 'foo' in tbl.columns, tbl.columns
def test_ensure_column(self):
tbl = self.tbl
tbl.create_column_by_example('foo', 0.1)
assert 'foo' in tbl.table.c, tbl.table.c
assert isinstance(tbl.table.c['foo'].type, FLOAT), \
tbl.table.c['bar'].type
tbl.create_column_by_example('bar', 1)
assert 'bar' in tbl.table.c, tbl.table.c
assert isinstance(tbl.table.c['bar'].type, BIGINT), \
tbl.table.c['bar'].type
tbl.create_column_by_example('pippo', 'test')
assert 'pippo' in tbl.table.c, tbl.table.c
assert isinstance(tbl.table.c['pippo'].type, TEXT), \
tbl.table.c['pippo'].type
tbl.create_column_by_example('bigbar', 11111111111)
assert 'bigbar' in tbl.table.c, tbl.table.c
assert isinstance(tbl.table.c['bigbar'].type, BIGINT), \
tbl.table.c['bigbar'].type
tbl.create_column_by_example('littlebar', -11111111111)
assert 'littlebar' in tbl.table.c, tbl.table.c
assert isinstance(tbl.table.c['littlebar'].type, BIGINT), \
tbl.table.c['littlebar'].type
def test_key_order(self):
res = self.db.query('SELECT temperature, place FROM weather LIMIT 1')
keys = list(res.next().keys())
assert keys[0] == 'temperature'
assert keys[1] == 'place'
def test_empty_query(self):
empty = list(self.tbl.find(place='not in data'))
assert len(empty) == 0, empty
class Constructor(dict):
""" Very simple low-functionality extension to ``dict`` to
provide attribute access to dictionary contents"""
def __getattr__(self, name):
return self[name]
class RowTypeTestCase(unittest.TestCase):
def setUp(self):
self.db = connect('sqlite:///:memory:', row_type=Constructor)
self.tbl = self.db['weather']
for row in TEST_DATA:
self.tbl.insert(row)
def tearDown(self):
for table in self.db.tables:
self.db[table].drop()
def test_find_one(self):
self.tbl.insert({
'date': datetime(2011, 1, 2),
'temperature': -10,
'place': 'Berlin'}
)
d = self.tbl.find_one(place='Berlin')
assert d['temperature'] == -10, d
assert d.temperature == -10, d
d = self.tbl.find_one(place='Atlantis')
assert d is None, d
def test_find(self):
ds = list(self.tbl.find(place=TEST_CITY_1))
assert len(ds) == 3, ds
for item in ds:
assert isinstance(item, Constructor), item
ds = list(self.tbl.find(place=TEST_CITY_1, _limit=2))
assert len(ds) == 2, ds
for item in ds:
assert isinstance(item, Constructor), item
def test_distinct(self):
x = list(self.tbl.distinct('place'))
assert len(x) == 2, x
for item in x:
assert isinstance(item, Constructor), item
x = list(self.tbl.distinct('place', 'date'))
assert len(x) == 6, x
for item in x:
assert isinstance(item, Constructor), item
def test_iter(self):
c = 0
for row in self.tbl:
c += 1
assert isinstance(row, Constructor), row
assert c == len(self.tbl)
if __name__ == '__main__':
unittest.main()