def _mock_allocated(self):
def mock_person(serialno, age, sex, income):
return {
'serial_number': serialno,
'age': age,
'sex': sex,
'individual_income': income
}
allocated_persons = pandas.DataFrame([
mock_person('b', '35-64', 'F', '40k+'),
mock_person('b', '35-64', 'M', 'None'),
])
def mock_household(serialno, num_people, num_vehicles, income, count,
tract):
return {
'serial_number': serialno,
'num_people': num_people,
'num_vehicles': num_vehicles,
'household_income': income,
'count': count,
'tract': tract,
}
allocated_households = pandas.DataFrame([
mock_household('b', '6+', '2', '40k+', count=2, tract='tract1'),
mock_household('b', '6+', '2', '40k+', count=2, tract='tract2'),
])
return HouseholdAllocator(allocated_households, allocated_persons)
def test_write_to_file(self):
persons = MagicMock()
households = MagicMock()
allocator = HouseholdAllocator(households, persons)
allocator.write(person_file='persons_file',
household_file='households_file')
persons.to_csv.assert_called_once_with('persons_file')
households.to_csv.assert_called_once_with('households_file')