diff --git a/source/jormungandr/jormungandr/scenarios/new_default.py b/source/jormungandr/jormungandr/scenarios/new_default.py index 4ac962740f..b32eb379f4 100644 --- a/source/jormungandr/jormungandr/scenarios/new_default.py +++ b/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) diff --git a/source/jormungandr/jormungandr/scenarios/tests/new_default_tests.py b/source/jormungandr/jormungandr/scenarios/tests/new_default_tests.py index 6e3985ccf4..d0f7ae8566 100644 --- a/source/jormungandr/jormungandr/scenarios/tests/new_default_tests.py +++ b/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)