Skip to content

Commit

Permalink
Add tests for the archive task
Browse files Browse the repository at this point in the history
  • Loading branch information
HeyZoos committed Feb 3, 2017
1 parent a0e1151 commit cbe4876
Show file tree
Hide file tree
Showing 4 changed files with 109 additions and 5 deletions.
2 changes: 2 additions & 0 deletions .travis.yml
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ addons:
packages:
- gdal-bin
- postgresql-9.4-postgis-2.3
- postgresql-9.4-plv8
services:
- postgresql
- redis-server
Expand Down Expand Up @@ -48,6 +49,7 @@ script:
- nosetests --nologcapture tests/submission/ -v
- nosetests --nologcapture tests/test_sensor_network/test_sensor_networks.py -v
- nosetests --nologcapture tests/test_sensor_network/test_nearest.py -v
- nosetests --nologcapture tests/test_tasks/test_archive.py -v

deploy:
- provider: elasticbeanstalk
Expand Down
10 changes: 5 additions & 5 deletions plenario/tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -173,25 +173,25 @@ def update_weather() -> True:


@worker.task()
def archive(path: str, table: str, start: str, end: str) -> bool:
def archive(path: str, table: str, start: str, end: str) -> str:
"""Store the results of a query in a tar.gz file."""

redshift_base.metadata.reflect()

table = redshift_base.metadata.tables[table]
query = redshift_session.query(table) \
.filter(table.c.datetime >= date_parse(start)) \
.filter(table.c.datetime < date_parse(end))
temp = tempfile.NamedTemporaryFile()
.filter(table.c.datetime <= date_parse(end))
temp = tempfile.NamedTemporaryFile('wt')
writer = csv.writer(temp)

for row in query.yield_per(1000):
writer.writerow(row)
temp.file.seek(0)

tar = tarfile.open(path, mode='w:gz')

tar.add(temp.name)
tar.close()
temp.close()

return True
return temp.name
Empty file added tests/test_tasks/__init__.py
Empty file.
102 changes: 102 additions & 0 deletions tests/test_tasks/test_archive.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,102 @@
import csv
import os
import tarfile
import unittest

from datetime import timedelta
from dateutil.parser import parse as date_parse
from sqlalchemy import create_engine, MetaData
from sqlalchemy import Table, Column, String, DateTime, Float
from sqlalchemy.exc import ProgrammingError
from sqlalchemy.orm import sessionmaker

from plenario.tasks import archive


uri = 'postgresql://postgres:password@localhost:5432'
engine = create_engine(uri)
connection = engine.connect()
metadata = MetaData()


temperature = Table(
'temperature',
metadata,
Column('node_id', String),
Column('datetime', DateTime),
Column('meta_id', Float),
Column('sensor', String))


class TestArchive(unittest.TestCase):

@classmethod
def setUpClass(cls):

connection.execute('commit')
try:
connection.execute('drop database plenario_test')
except ProgrammingError:
pass
connection.execute('commit')
connection.execute('create database plenario_test')

cls.engine = create_engine(uri + '/plenario_test')
cls.session = sessionmaker(bind=cls.engine)()
temperature.create(cls.engine)

time = date_parse('2016-01-01 00:00:00')
for i in range(0, 100):
insert = temperature.insert().values(
node_id='0',
datetime=time + timedelta(days=i),
meta_id='0',
sensor='0')
cls.session.execute(insert)
cls.session.commit()

def test_archive_january(self):

name = archive('./t.tar.gz', 'temperature', '2016-01-01', '2016-01-31')

tar = tarfile.open('t.tar.gz')
tar.extractall()
tar.close()

filename = name.split('/', 1)[-1]
file = open(filename)
reader = csv.reader(file)

line_num = 0
for row in reader:
line_num += 1

self.assertEqual(line_num, 31)

file.close()
os.remove(filename)
os.rmdir('tmp')
os.remove('t.tar.gz')

def test_archive_february(self):

name = archive('./t.tar.gz', 'temperature', '2016-02-01', '2016-02-29')

tar = tarfile.open('t.tar.gz')
tar.extractall()
tar.close()

filename = name.split('/', 1)[-1]
file = open(filename)
reader = csv.reader(file)

line_num = 0
for row in reader:
line_num += 1

self.assertEqual(line_num, 29)

file.close()
os.remove(filename)
os.rmdir('tmp')
os.remove('t.tar.gz')

0 comments on commit cbe4876

Please sign in to comment.