Merge pull request #4230 from hove-io/keep_best_olympics_when_using_max_nb_journeys

[Jormungandr] Keep best_olympics

Author: xlqian

source/jormungandr/jormungandr/scenarios/new_default.py

@@ -101,7 +101,9 @@
 from functools import cmp_to_key
 
 SECTION_TYPES_TO_RETAIN = {response_pb2.PUBLIC_TRANSPORT, response_pb2.STREET_NETWORK}
+JOURNEY_TAGS_TO_RETAIN = ['best_olympics']
 JOURNEY_TYPES_TO_RETAIN = ['best', 'comfort', 'non_pt_walk', 'non_pt_bike', 'non_pt_bss']
+JOURNEY_TYPES_SCORE = {t: i for i, t in enumerate(JOURNEY_TYPES_TO_RETAIN)}
 STREET_NETWORK_MODE_TO_RETAIN = {response_pb2.Ridesharing, response_pb2.Car, response_pb2.Bike, response_pb2.Bss}
 TEMPLATE_MSG_UNKNOWN_OBJECT = "The entry point: {} is not valid"
 
@@ -515,7 +517,7 @@ def _build_candidate_pool_and_sections_set(journeys):
     idx_of_jrny_must_keep = list()
 
     for (i, jrny) in enumerate(journeys):
-        if jrny.type in JOURNEY_TYPES_TO_RETAIN:
+        if set(jrny.tags) & set(JOURNEY_TAGS_TO_RETAIN) or jrny.type in set(JOURNEY_TYPES_TO_RETAIN):
             idx_of_jrny_must_keep.append(i)
         sections_set |= set([_get_section_id(s) for s in jrny.sections if s.type in SECTION_TYPES_TO_RETAIN])
         candidates_pool.append(jrny)
@@ -750,16 +752,16 @@ def _inverse_selection(d, indexes):
         )
 
         # At this point, resp.journeys should contain only must-have journeys
-        list_dict = collections.defaultdict(list)
-        for jrny in resp.journeys:
-            if not journey_filter.to_be_deleted(jrny):
-                list_dict[jrny.type].append(jrny)
+        must_have = [j for j in resp.journeys if not journey_filter.to_be_deleted(j)]
 
-        sorted_by_type_journeys = []
-        for t in JOURNEY_TYPES_TO_RETAIN:
-            sorted_by_type_journeys.extend(list_dict.get(t, []))
+        def journey_score(j):
+            if set(j.tags) & set(JOURNEY_TAGS_TO_RETAIN):
+                return -100
+            return JOURNEY_TYPES_SCORE.get(j.type, float('inf'))
 
-        for jrny in sorted_by_type_journeys[max_nb_journeys:]:
+        must_have.sort(key=journey_score)
+
+        for jrny in must_have[max_nb_journeys:]:
             journey_filter.mark_as_dead(jrny, is_debug, 'Filtered by max_nb_journeys')
 
         journey_filter.delete_journeys((resp,), request)

source/jormungandr/jormungandr/scenarios/tests/new_default_tests.py

@@ -282,6 +282,36 @@ def culling_journeys_4_test():
         assert jrny.type in ('best', 'comfort', 'non_pt_walk')
 
 
+def culling_journeys_5_test():
+    """
+    When max_nb_journeys == 3 and nb_must_have_journeys == 4 and one journey has best_olympics in its tags
+    """
+    mocked_pb_response = build_mocked_response()
+    # tag last journeys with
+    mocked_pb_response.journeys[10].tags.append("best_olympics")
+    mocked_request = {'max_nb_journeys': 6, 'debug': False, 'datetime': 1444903200}
+    new_default.culling_journeys(mocked_pb_response, mocked_request)
+
+    best_olympic = next((j for j in mocked_pb_response.journeys if 'best_olympics' in j.tags))
+    assert best_olympic
+
+
+def culling_journeys_6_test():
+    """
+    When max_nb_journeys == 3 and nb_must_have_journeys == 4 and one journey has best_olympics in its tags
+    """
+    mocked_pb_response = build_mocked_response()
+    # tag last journeys with
+    mocked_pb_response.journeys[5].tags.append("best_olympics")
+
+    mocked_request = {'max_nb_journeys': 3, 'debug': False, 'datetime': 1444903200}
+    new_default.culling_journeys(mocked_pb_response, mocked_request)
+    assert len(mocked_pb_response.journeys) == 3
+
+    best_olympic = next((j for j in mocked_pb_response.journeys if 'best_olympics' in j.tags))
+    assert best_olympic
+
+
 def aggregate_journeys_test():
     mocked_pb_response = build_mocked_response()
     aggregated_journeys, remaining_journeys = new_default.aggregate_journeys(mocked_pb_response.journeys)