def test_transpose(self):
new_tbl = self.tbl.transpose(include_header=False)
self.assertEqual(
new_tbl,
pgreaper.Table(name='Countries',
row_values=[['Washington', 'Moscow', 'Ottawa'],
['USA', 'Russia', 'Canada'],
['USD', 'RUB', 'CAD'],
['American', 'Russian', 'Canadian'],
[324774000, 144554993, 35151728]]))
def test_from_nothing(self):
''' Test that adding to a Table with no columns or rows works '''
table = pgreaper.Table(name=None)
table.add_dicts([{
"Capital": "Beijing",
"Country": "China",
"Demonym": 'Chinese',
"Population": 1373541278,
}])
self.assertEqual(table['Capital'], ['Beijing'])
def load_test3():
load_func = partial(json.loads)
file = open('persons_nested.json', mode='rb')
x = JSONStreamingDecoder(source=file,
loads=load_func)
tbl = pgreaper.Table(name='test',
col_names=['full_name', 'occupation', 'nationality'])
for i in x:
tbl.add_dict(i)
return tbl
def load_test():
load_func = partial(json.loads,
object_hook=lambda x: [v for k, v in \
x.items() if k in ['full_name', 'occupation', 'nationality']])
file = open('persons_nested.json', mode='rb')
x = JSONStreamingDecoder(source=file,
loads=load_func)
try:
rows = [i for i in x]
except:
x.source.close()
tbl = pgreaper.Table(name='test',
col_names=['full_name', 'occupation', 'nationality'],
row_values=rows)
return tbl
class UpsertTest(PostgresTestCase):
''' Test various UPSERT options '''
data = pgreaper.Table('Countries',
col_names=world_countries_cols(),
row_values=world_countries())
# p_key = 1,
data.append(["Beijing", "China", "CNY", 'Chinese', 1373541278])
data.p_key = 'Country'
drop_tables = ['countries']
def setUp(self):
super(UpsertTest, self).setUp()
# DROP TABLE IF EXISTS causes psycopg2 to hang...
try:
self.cursor.execute('DROP TABLE countries')
self.conn.commit()
except psycopg2.ProgrammingError:
self.conn.rollback()
def test_on_conflict_do_nothing(self):
pgreaper.table_to_pg(self.data[0:2], name='countries', dbname=TEST_DB)
assert (self.data[0:2].p_key == 1)
pgreaper.table_to_pg(self.data[2:],
name='countries',
dbname=TEST_DB,
on_p_key='nothing')
self.assertCount('countries', 4)
def test_insert_or_replace(self):
pgreaper.table_to_pg(UpsertTest.data, name='countries', dbname=TEST_DB)
# Set entire population column to 0
self.data.apply('Population', lambda x: 0)
pgreaper.table_to_pg(self.data,
on_p_key='replace',
name='countries',
dbname=TEST_DB)
self.cursor.execute('SELECT sum(population::bigint) FROM countries')
self.assertEqual(self.cursor.fetchall()[0][0], 0)
def test_reorder_do_nothing(self):
''' Test if SQLify can reorder input to match SQL table '''
# Load Table
pgreaper.table_to_pg(UpsertTest.data[0:2],
name='countries',
dbname=TEST_DB)
# Make Table with same columns, different order then try to load
data = self.data[2:].reorder(4, 3, 2, 0, 1)
pgreaper.table_to_pg(data,
name='countries',
dbname=TEST_DB,
reorder=True,
on_p_key='nothing')
self.assertCount('countries', 4)
def test_append(self):
''' Regular append, no primary key constraint '''
self.data.p_key = None
pgreaper.table_to_pg(self.data, name='countries', dbname=TEST_DB)
pgreaper.table_to_pg(self.data,
name='countries',
dbname=TEST_DB,
append=True)
self.assertCount('countries', 8)
def test_expand_input(self, expand_input=True):
''' Test that input expansion is handled properly '''
pgreaper.table_to_pg(self.data[0:2], name='countries', dbname=TEST_DB)
# Needs to line up columns correctly with existing schema
needs_expanding = self.data[2:].subset('Country', 'Population')
pgreaper.table_to_pg(needs_expanding,
name='countries',
dbname=TEST_DB,
expand_input=expand_input)
# Check that expanded columns were filled with NULLs
self.cursor.execute('SELECT count(*) FROM countries WHERE'
' capital is NULL')
count = self.cursor.fetchall()[0][0]
self.assertEqual(count, 2)
def test_no_expand_input(self):
''' Test that input expansion doesn't happen unless explicitly specified '''
with self.assertRaises(ValueError):
self.test_expand_input(expand_input=False)
def test_expand_sql(self, expand_sql=True):
''' Test that adding columns to SQL tables works '''
data = copy.deepcopy(self.data)
data.delete('Population')
# Upload truncated data set
pgreaper.table_to_pg(data[0:2], name='countries', dbname=TEST_DB)
pgreaper.table_to_pg(self.data[2:],
name='countries',
dbname=TEST_DB,
expand_sql=expand_sql)
# Test that extra columns were added
self.cursor.execute('SELECT count(population) FROM countries '
'WHERE population IS NOT NULL')
self.assertEqual(self.cursor.fetchall()[0][0], 2)
def test_no_expand_sql(self):
'''
Test that adding columns to SQL tables
doesn't happen unless explicitly specified
'''
with self.assertRaises(ValueError):
self.test_expand_sql(expand_sql=False)