diff --git a/src/google_flight_analysis/flight.py b/src/google_flight_analysis/flight.py index 6bef828..a7e0373 100644 --- a/src/google_flight_analysis/flight.py +++ b/src/google_flight_analysis/flight.py @@ -105,7 +105,7 @@ def time_arrive(self): return self._time_arrive def _classify_arg(self, arg): - if ('AM' in arg or 'PM' in arg) and len(self._times) < 2: + if ('AM' in arg or 'PM' in arg) and len(self._times) < 2 and ':' in arg: # arrival or departure time delta = timedelta(days = 0) if arg[-2] == '+': diff --git a/src/google_flight_analysis/scrape.py b/src/google_flight_analysis/scrape.py index 9dbc7a8..ad2b8db 100644 --- a/src/google_flight_analysis/scrape.py +++ b/src/google_flight_analysis/scrape.py @@ -146,8 +146,8 @@ def _set_properties(self, *args): self._url = self._make_url() self._type = 'round-trip' - # chain-trip - elif len(args) >= 3 and len(args) % 3 == 0: + # chain-trip, chain is component of 3s, check that last one is an actual date to not confuse w perfect + elif len(args) >= 3 and len(args) % 3 == 0 and len(args[-1]) == 10 and type(args[-1]) == str: self._origin, self._dest, self._date = [], [], [] for i in range(0, len(args), 3): @@ -173,7 +173,7 @@ def _set_properties(self, *args): for i in range(2, len(args)-1, 2): assert len(args[i]) == 3 and type(args[i]) == str, "Issue with arg {}, see docs".format(i) - assert len(args[i + 1] == 10 and type(args[i + 1])) == str, "Issue with arg {}, see docs".format(i+1) + assert len(args[i + 1]) == 10 and type(args[i + 1]) == str, "Issue with arg {}, see docs".format(i+1) self._origin += [args[i]] self._dest += [args[i]] diff --git a/tests/test_class.py b/tests/test_class.py index 26d06b6..70c748a 100644 --- a/tests/test_class.py +++ b/tests/test_class.py @@ -30,7 +30,7 @@ def func_0(): res5 = Scrape("JFK", "AMS", "2023-11-10", "CDG", "AMS", "2023-11-17", "AMS", "IST", "2023-11-25") # perfect chain -res6 = Scrape("JFK", "2023-11-10", "AMS", "2023-11-17", "CDG", "2023-11-17", "IST", "2023-11-25") +res6 = Scrape("JFK", "2023-11-10", "AMS", "2023-11-17", "CDG", "2023-11-20", "IST", "2023-11-25", "JFK") '''os.system('rm tests/test_data/LGA-RDU.csv') os.system('rm tests/test_data/CDG-IST.csv') @@ -112,13 +112,13 @@ def test_19(): #-------QUERY 6 def test_20(): - assert res6.origin == ["JFK", "AMS", "CDG"], "Test 20 Failed." + assert res6.origin == ["JFK", "AMS", "CDG", "IST"], "Test 20 Failed." def test_21(): - assert res6.dest == ["AMS", "CDG", "IST"], "Test 21 Failed." + assert res6.dest == ["AMS", "CDG", "IST", "JFK"], "Test 21 Failed." def test_22(): - assert res6.date == ["2023-11-10", "2023-11-17", "2023-11-25"], "Test 22 Failed." + assert res6.date == ["2023-11-10", "2023-11-17", "2023-11-20", "2023-11-25"], "Test 22 Failed." '''#-------CACHE 1