|
18
|
1 #!/usr/bin/env python
|
|
|
2
|
|
|
3 import csv
|
|
|
4 import os
|
|
|
5 import shutil
|
|
|
6 import tempfile
|
|
|
7 import unittest
|
|
|
8 from collections import OrderedDict
|
|
|
9 from orion.columns import read_columns
|
|
|
10 from orion.transpose import transpose
|
|
|
11 from StringIO import StringIO
|
|
|
12
|
|
|
13
|
|
|
14 class TestColumns(unittest.TestCase):
|
|
|
15
|
|
|
16 # fake data
|
|
|
17 header = ["Animal", "Baby", "Group"]
|
|
|
18 data = [("Fox", "Kit", "Leash"),
|
|
|
19 ("Whale", "Calf", "Pod"),
|
|
|
20 ("Cuttlefish", "Hatchling", "Bob"),
|
|
|
21 ("Crow", "Hatchling", "Murder")]
|
|
|
22
|
|
|
23 def create_data(self, fp):
|
|
|
24 """write simple data to a test file and return its path"""
|
|
|
25
|
|
|
26 writer = csv.writer(fp)
|
|
|
27 writer.writerow(self.header)
|
|
|
28 writer.writerows(self.data)
|
|
|
29 fp.flush()
|
|
|
30
|
|
|
31 def fake_data_assertions(self, columns):
|
|
|
32 """test the columns created from our fake data"""
|
|
|
33
|
|
|
34 assert len(columns) == len(self.header)
|
|
|
35 assert set(columns.keys()) == set(self.header)
|
|
|
36 transposed = transpose(self.data)
|
|
|
37 assert set([len(row) for row in transposed]) == set([len(self.data)])
|
|
|
38 for key, column in zip(self.header, transposed):
|
|
|
39 assert len(columns[key]) == len(self.data)
|
|
|
40 assert columns[key] == column
|
|
|
41
|
|
|
42 def test_string(self):
|
|
|
43 """basic test of columns interface"""
|
|
|
44
|
|
|
45 # write fake data
|
|
|
46 fileobj = StringIO()
|
|
|
47 self.create_data(fileobj)
|
|
|
48 fileobj.seek(0)
|
|
|
49
|
|
|
50 # read into columns
|
|
|
51 columns = read_columns(fileobj)
|
|
|
52
|
|
|
53 # test what we have
|
|
|
54 self.fake_data_assertions(columns)
|
|
|
55
|
|
|
56 def test_file(self):
|
|
|
57 """write to a real file and use that"""
|
|
|
58
|
|
|
59 tmpdir = tempfile.mkdtemp()
|
|
|
60 try:
|
|
|
61 # write the data
|
|
|
62 path = os.path.join(tmpdir, 'tmp.csv')
|
|
|
63 with open(path, 'w') as f:
|
|
|
64 self.create_data(f)
|
|
|
65 assert os.path.exists(path)
|
|
|
66
|
|
|
67 # read into columns
|
|
|
68 columns = read_columns(path)
|
|
|
69
|
|
|
70 # test what we have
|
|
|
71 self.fake_data_assertions(columns)
|
|
|
72 finally:
|
|
|
73 shutil.rmtree(tmpdir, ignore_errors=True)
|
|
|
74
|
|
|
75 def test_ordered_dict(self):
|
|
|
76 """ensure we can make use of an OrderedDict"""
|
|
|
77
|
|
|
78 # write fake data
|
|
|
79 fileobj = StringIO()
|
|
|
80 self.create_data(fileobj)
|
|
|
81 fileobj.seek(0)
|
|
|
82
|
|
|
83 # read into columns
|
|
|
84 columns = read_columns(fileobj, type=OrderedDict)
|
|
|
85
|
|
|
86 # test what we have
|
|
|
87 self.fake_data_assertions(columns)
|
|
|
88
|
|
|
89
|
|
|
90 if __name__ == '__main__':
|
|
|
91 unittest.main()
|
