From e0240130935d8748b14ed31f10b636daa615316a Mon Sep 17 00:00:00 2001
From: zhexu14 <64713351+zhexu14@users.noreply.github.com>
Date: Mon, 29 May 2023 23:25:29 +1000
Subject: [PATCH] issue 2922: make BAI plannable against missile and costal
 sites

---
 game/theater/theatergroundobject.py       | 16 ++++++++++++++++
 tests/theater/test_theatergroundobject.py |  3 ++-
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/game/theater/theatergroundobject.py b/game/theater/theatergroundobject.py
index e302c94b..28834a9d 100644
--- a/game/theater/theatergroundobject.py
+++ b/game/theater/theatergroundobject.py
@@ -434,6 +434,14 @@ class MissileSiteGroundObject(TheaterGroundObject):
     def should_head_to_conflict(self) -> bool:
         return True
 
+    def mission_types(self, for_player: bool) -> Iterator[FlightType]:
+        from game.ato import FlightType
+
+        if not self.is_friendly(for_player):
+            yield FlightType.BAI
+        for mission_type in super().mission_types(for_player):
+            yield mission_type
+
 
 class CoastalSiteGroundObject(TheaterGroundObject):
     def __init__(
@@ -466,6 +474,14 @@ class CoastalSiteGroundObject(TheaterGroundObject):
     def should_head_to_conflict(self) -> bool:
         return True
 
+    def mission_types(self, for_player: bool) -> Iterator[FlightType]:
+        from game.ato import FlightType
+
+        if not self.is_friendly(for_player):
+            yield FlightType.BAI
+        for mission_type in super().mission_types(for_player):
+            yield mission_type
+
 
 class IadsGroundObject(TheaterGroundObject, ABC):
     def __init__(
diff --git a/tests/theater/test_theatergroundobject.py b/tests/theater/test_theatergroundobject.py
index 9de17d90..eabe72a7 100644
--- a/tests/theater/test_theatergroundobject.py
+++ b/tests/theater/test_theatergroundobject.py
@@ -180,7 +180,8 @@ def test_mission_types_enemy(mocker: Any) -> None:
             control_point=dummy_control_point,
         )
         mission_types = list(ground_object.mission_types(for_player=False))
-        assert len(mission_types) == 6
+        assert len(mission_types) == 7
+        assert FlightType.BAI in mission_types
         assert FlightType.STRIKE in mission_types
         assert FlightType.REFUELING in mission_types
         assert FlightType.ESCORT in mission_types