def test_country_flight_files(self):
# process_airport_file
airport1 = airports.Airport()
with open(self.airports_file, "r", encoding="UTF-8") as f:
reader = csv.reader(f)
for row in reader:
airport1.process_airport_row(row)
# process_flight_file
flight_per_country = country_flights.FlightPerCountry(
airport1.airports_by_iata)
with open(self.flights_file, "r", encoding="UTF-8") as f:
reader = csv.reader(f)
for row in reader:
flight_per_country.process_flight(row)
expected = {
"Afghanistan": [14, 29],
"Congo (Brazzaville)": [10, 34],
"Denmark": [23, 286],
"Guinea-Bissau": [0, 7],
"Sweden": [167, 307],
flight_per_country.unknown_country: [93, 636],
}
result = {}
for country in expected:
result[country] = flight_per_country.countries[country]
self.assertEqual(result, expected)
def test_get_airport_country_by_iata(self):
flightPerCountry = country_flights.FlightPerCountry(self.IataCountry)
expected = self.IataCountry
result = {}
for iata in self.IataCountry:
result[iata] = flightPerCountry.get_airport_country(
airport_code=iata
)
self.assertEqual(result, expected)
def main():
# go to src dir
src_dir = os.path.dirname(os.path.realpath(__file__))
# declare log file
logging.basicConfig(
filename=os.path.join(src_dir, "log", "flights.log"),
format="%(levelname)s: %(asctime)s: %(process)d: %(filename)s:"
" %(funcName)s: %(message)s",
level=logging.INFO,
)
logging.info("Process started!")
# output file
parser = argparse.ArgumentParser()
parser.add_argument(
"--output_file",
help="output file where results are stored",
default="output_data/output.csv",
)
args = parser.parse_args()
# input_data dir
data_dir = os.path.join(src_dir, "..", "input_data")
# process airports
airports_file = os.path.join(data_dir, "airports.dat")
airport1 = airports.Airport()
process_airport_file(airports_file, airport1)
# process flights
flights_file = os.path.join(data_dir, "routes.dat")
flight_per_country = country_flights.FlightPerCountry(
airport1.airports_by_iata)
process_flight_file(flights_file, flight_per_country)
# save results
results = flight_per_country.get_results_format1()
save_data(args.output_file, results)
msg = "Process completed!"
logging.info(msg)
print(msg)
def test_result_format(self):
"""this can check full results, order important
results are list of list.
"""
flightPerCountry = country_flights.FlightPerCountry(self.IataCountry)
expected = [
["Sweden", 1, 1],
["United States", 1, 1],
[flightPerCountry.unknown_country, 1, 2],
]
for row in self.FlightDataA:
flightPerCountry.process_flight(row)
result = flightPerCountry.get_results_format1()
self.assertEqual(result, expected)
def test_get_country_flights(self):
"""this can check partial results, order not important
for example United States is exluded in the test
but is pressent in the flightPerCountry.countries
"""
flightPerCountry = country_flights.FlightPerCountry(self.IataCountry)
expected = {
flightPerCountry.unknown_country: [1, 2],
# "United States": [1, 1],
"Sweden": [1, 1],
}
for row in self.FlightDataA:
flightPerCountry.process_flight(row)
result = {}
for country in expected:
result[country] = flightPerCountry.countries[country]
self.assertEqual(result, expected)
def test_get_unknown_country(self):
flightPerCountry = country_flights.FlightPerCountry({})
expected = flightPerCountry.unknown_country
result = flightPerCountry.get_airport_country("KIO")
self.assertEqual(result, expected)