diff --git a/game/ato/flightplans/cas.py b/game/ato/flightplans/cas.py index 5a379ced..ef8472e2 100644 --- a/game/ato/flightplans/cas.py +++ b/game/ato/flightplans/cas.py @@ -63,12 +63,6 @@ class CasFlightPlan(PatrollingFlightPlan[CasLayout], UiZoneDisplay): def combat_speed_waypoints(self) -> set[FlightWaypoint]: return {self.layout.ingress, self.layout.patrol_start, self.layout.patrol_end} - def request_escort_at(self) -> FlightWaypoint | None: - return self.layout.patrol_start - - def dismiss_escort_at(self) -> FlightWaypoint | None: - return self.layout.patrol_end - def ui_zone(self) -> UiZone: midpoint = ( self.layout.patrol_start.position + self.layout.patrol_end.position @@ -128,6 +122,7 @@ class Builder(IBuilder[CasFlightPlan, CasLayout]): patrol_start_waypoint.name = "FLOT START" patrol_start_waypoint.pretty_name = "FLOT start" patrol_start_waypoint.description = "FLOT boundary" + patrol_start_waypoint.wants_escort = True patrol_end_waypoint = builder.nav( patrol_end, patrol_altitude, use_agl_patrol_altitude @@ -135,6 +130,7 @@ class Builder(IBuilder[CasFlightPlan, CasLayout]): patrol_end_waypoint.name = "FLOT END" patrol_end_waypoint.pretty_name = "FLOT end" patrol_end_waypoint.description = "FLOT boundary" + patrol_end_waypoint.wants_escort = True ingress = builder.ingress( FlightWaypointType.INGRESS_CAS, ingress_point, location diff --git a/game/ato/flightplans/flightplan.py b/game/ato/flightplans/flightplan.py index b38682d8..5c1e1618 100644 --- a/game/ato/flightplans/flightplan.py +++ b/game/ato/flightplans/flightplan.py @@ -205,24 +205,21 @@ class FlightPlan(ABC, Generic[LayoutT]): raise NotImplementedError def request_escort_at(self) -> FlightWaypoint | None: - return None + try: + return next(self.escorted_waypoints()) + except StopIteration: + return None def dismiss_escort_at(self) -> FlightWaypoint | None: - return None + try: + return list(self.escorted_waypoints())[-1] + except IndexError: + return None def escorted_waypoints(self) -> Iterator[FlightWaypoint]: - begin = self.request_escort_at() - end = self.dismiss_escort_at() - if begin is None or end is None: - return - escorting = False - for waypoint in self.waypoints: - if waypoint == begin: - escorting = True - if escorting: + for waypoint in self.iter_waypoints(): + if waypoint.wants_escort: yield waypoint - if waypoint == end: - return def takeoff_time(self) -> datetime: return self.tot - self._travel_time_to_waypoint(self.tot_waypoint) diff --git a/game/ato/flightplans/formation.py b/game/ato/flightplans/formation.py index 81bd7290..4ba8ab1c 100644 --- a/game/ato/flightplans/formation.py +++ b/game/ato/flightplans/formation.py @@ -37,12 +37,6 @@ class FormationFlightPlan(LoiterFlightPlan[LayoutT], ABC): def combat_speed_waypoints(self) -> set[FlightWaypoint]: return self.package_speed_waypoints - def request_escort_at(self) -> FlightWaypoint | None: - return self.layout.join - - def dismiss_escort_at(self) -> FlightWaypoint | None: - return self.layout.split - @cached_property def best_flight_formation_speed(self) -> Speed: """The best speed this flight is capable at all formation waypoints. diff --git a/game/ato/flightplans/formationattack.py b/game/ato/flightplans/formationattack.py index def0b6c4..0e9e0f2a 100644 --- a/game/ato/flightplans/formationattack.py +++ b/game/ato/flightplans/formationattack.py @@ -145,7 +145,18 @@ class FormationAttackBuilder(IBuilder[FlightPlanT, LayoutT], ABC): hold = builder.hold(self._hold_point()) join = builder.join(self.package.waypoints.join) + join.wants_escort = True + + ingress = builder.ingress( + ingress_type, self.package.waypoints.ingress, self.package.target + ) + ingress.wants_escort = True + + for target_waypoint in target_waypoints: + target_waypoint.wants_escort = True + split = builder.split(self.package.waypoints.split) + split.wants_escort = True refuel = builder.refuel(self.package.waypoints.refuel) return FormationAttackLayout( @@ -155,9 +166,7 @@ class FormationAttackBuilder(IBuilder[FlightPlanT, LayoutT], ABC): hold.position, join.position, self.doctrine.ingress_altitude ), join=join, - ingress=builder.ingress( - ingress_type, self.package.waypoints.ingress, self.package.target - ), + ingress=ingress, targets=target_waypoints, split=split, refuel=refuel, diff --git a/game/ato/flightwaypoint.py b/game/ato/flightwaypoint.py index b261f579..590d0745 100644 --- a/game/ato/flightwaypoint.py +++ b/game/ato/flightwaypoint.py @@ -41,6 +41,8 @@ class FlightWaypoint: # The minimum amount of fuel remaining at this waypoint in pounds. min_fuel: float | None = None + wants_escort: bool = False + actions: list[WaypointAction] = field(default_factory=list) options: dict[str, WaypointOption] = field(default_factory=dict) diff --git a/tests/ato/test_flightplan.py b/tests/ato/test_flightplan.py new file mode 100644 index 00000000..ed35683f --- /dev/null +++ b/tests/ato/test_flightplan.py @@ -0,0 +1,85 @@ +from typing import cast, Any + +import pytest +from dcs import Point +from dcs.terrain import Caucasus + +from game.ato import Flight, FlightWaypoint +from game.ato.flightplans.custom import CustomFlightPlan, CustomLayout +from game.ato.flightplans.flightplan import FlightPlan +from game.ato.flightwaypointtype import FlightWaypointType + + +@pytest.fixture(name="unescorted_flight_plan") +def unescorted_flight_plan_fixture() -> FlightPlan[Any]: + point = Point(0, 0, Caucasus()) + departure = FlightWaypoint("", FlightWaypointType.TAKEOFF, point) + + waypoints = [ + FlightWaypoint(f"{i}", FlightWaypointType.NAV, point) for i in range(10) + ] + return CustomFlightPlan(cast(Flight, object()), CustomLayout(departure, waypoints)) + + +@pytest.fixture(name="escorted_flight_plan") +def escorted_flight_plan_fixture() -> FlightPlan[Any]: + point = Point(0, 0, Caucasus()) + departure = FlightWaypoint("", FlightWaypointType.TAKEOFF, point) + + waypoints = [ + FlightWaypoint(f"{i}", FlightWaypointType.NAV, point) for i in range(10) + ] + waypoints[1].wants_escort = True + waypoints[2].wants_escort = True + waypoints[3].wants_escort = True + waypoints[5].wants_escort = True + waypoints[7].wants_escort = True + waypoints[8].wants_escort = True + return CustomFlightPlan(cast(Flight, object()), CustomLayout(departure, waypoints)) + + +def test_escorted_flight_plan_escorted_waypoints( + escorted_flight_plan: FlightPlan[Any], +) -> None: + assert [w.name for w in escorted_flight_plan.escorted_waypoints()] == [ + "1", + "2", + "3", + "5", + "7", + "8", + ] + + +def test_escorted_flight_plan_request_escort_at( + escorted_flight_plan: FlightPlan[Any], +) -> None: + wp = escorted_flight_plan.request_escort_at() + assert wp is not None + assert wp.name == "1" + + +def test_escorted_flight_plan_dismiss_escort_at( + escorted_flight_plan: FlightPlan[Any], +) -> None: + wp = escorted_flight_plan.dismiss_escort_at() + assert wp is not None + assert wp.name == "8" + + +def test_unescorted_flight_plan_escorted_waypoints( + unescorted_flight_plan: FlightPlan[Any], +) -> None: + assert not list(unescorted_flight_plan.escorted_waypoints()) + + +def test_unescorted_flight_plan_request_escort_at( + unescorted_flight_plan: FlightPlan[Any], +) -> None: + assert unescorted_flight_plan.request_escort_at() is None + + +def test_unescorted_flight_plan_dismiss_escort_at( + unescorted_flight_plan: FlightPlan[Any], +) -> None: + assert unescorted_flight_plan.dismiss_escort_at() is None