def test_writerows(self):
'''
Write rows to a temporary file with rules.
'''
with tempfile.TemporaryFile(mode='w') as csvfile:
r = rigidity.Rigidity(csv.writer(csvfile))
r.writerows([['a', 'b'], ['c', 'd']])
def test_validate_write_no_rules(self):
'''
Test that the validate_write() method does not change the data
if no rules are provided.
'''
r = rigidity.Rigidity(csv.reader(tempfile.TemporaryFile()), [[], []])
data1 = [['a', 'b'], ['c', 'd']]
self.assertEqual(r.validate_write(data1), data1)
def test_writeheader(self):
'''
Test that the writehader() method on a CSVWriter is called when
Rigidity's csvheader() is called.
'''
writer = mock.MagicMock()
r = rigidity.Rigidity(writer, [])
r.writeheader()
writer.writeheader.assert_called_once_with()
def test_validate_dict_write_no_rules(self):
'''
Test that the validate_write() method works correctly on Dict
rows when the rules Dict is empty.
'''
r = rigidity.Rigidity(csv.reader(tempfile.TemporaryFile()), {
'a': [],
'b': []
})
data1 = {'a': 'hello', 'b': 'world'}
self.assertDictEqual(r.validate_write(data1), data1)
def test_writerow(self):
'''
Test that the writerow method calls the writerow method of the
CSVWriter object with validated/corrected data.
'''
# Test without rules
writer = mock.MagicMock()
r = rigidity.Rigidity(writer, [[], []])
r.writerow(('hello', 'world'))
writer.writerow.assert_called_with(['hello', 'world'])
r.writerow([1, 2])
writer.writerow.assert_called_with([1, 2])
# Test with rules
writer = mock.MagicMock()
r = rigidity.Rigidity(
writer, [[rigidity.rules.Drop()], [rigidity.rules.Drop()]])
r.writerow(('hello', 'world'))
writer.writerow.assert_called_with(['', ''])
r.writerow([1, 2])
writer.writerow.assert_called_with(['', ''])
def test___next__(self):
'''
Test that the __next__ method returns the next row,
post-validation and post-correction.
'''
DATA_FILE = 'data_0001.csv'
rules = [[rigidity.rules.Strip()], [rigidity.rules.Strip()]]
with open(os.path.join(DATA_DIR, DATA_FILE)) as csvfile:
r = rigidity.Rigidity(csv.reader(csvfile), rules)
self.assertEqual(next(r), ['hello', 'world'])
self.assertEqual(next(r), ['things', 'great'])
def test_data_simple_read(self):
'''
Perform a read of an actual CSV file with rules.
'''
DATA_FILE = 'data_0001.csv'
rules = [[rigidity.rules.Strip()], [rigidity.rules.Strip()]]
with open(os.path.join(DATA_DIR, DATA_FILE)) as csvfile:
r = rigidity.Rigidity(csv.reader(csvfile), rules)
for row in r:
self.assertEqual(row[0], row[0].strip())
self.assertEqual(row[1], row[1].strip())
def test___delattr___reader_attribute(self):
'''
Test that the __delattr__ method will attempt to delete from
the CSVReader or CSVWriter object when it has an attribute.
In this case, the deletion should cause an AttributeError
because the attribute cannot be deleted.
'''
DATA_FILE = 'data_0001.csv'
with open(os.path.join(DATA_DIR, DATA_FILE)) as csvfile:
reader = csv.DictReader(csvfile, fieldnames=['a', 'b'])
assert hasattr(reader, 'fieldnames')
r = rigidity.Rigidity(reader)
self.assertRaises(AttributeError, delattr, r, 'fieldnames')
'''
:returns: A dict containing the data stored in this object.
'''
return {
'name': self.name,
'website': self.website,
'description': self.description,
'attributes': self.attributes.base64().decode('ascii'),
'tables': self.tables
}
if __name__ == '__main__':
reader = csv.DictReader(open('dataset.csv'))
r = rigidity.Rigidity(reader,
rules,
display=rigidity.Rigidity.DISPLAY_SIMPLE)
r.skip() # Skip header
company_id = 1
companies = {}
for row in r:
print(row)
company = Company(company_id, row['name'], row['website'],
utils.truncate_description(row['description']))
# Iterate over all majors and add their data to the company
for major in MAJORS:
column_heading = constants.MAJOR_SPREDSHEET_HEADIGNS[major]
position_tuple = utils.extract_position_tuple(row[column_heading])
company.add_major(major, *position_tuple)
def test_writerow(self):
writer = mock.MagicMock()
r_rules = [[rules.Integer(action=rules.Integer.ACTION_DROPROW)]]
r = rigidity.Rigidity(writer, r_rules)
r.writerow(['a'])
self.assertFalse(writer.writerow.called)
def test___delattr___invalid_attribute(self):
r = rigidity.Rigidity(None, [[], []])
self.assertRaises(AttributeError, delattr, r, 'does_not_exist')