From a54f836cdecbdd9b893345e35896d316eacdda7c 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 2a9f0870..9d4046f8 100644
--- a/game/theater/theatergroundobject.py
+++ b/game/theater/theatergroundobject.py
@@ -448,6 +448,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__(
@@ -480,6 +488,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 35e81d68..009dd0e6 100644
--- a/tests/theater/test_theatergroundobject.py
+++ b/tests/theater/test_theatergroundobject.py
@@ -181,7 +181,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