Merge pull request #3851 from Sleet01/Fix_3848_unable_to_assign_asf_t…

…o_transport Fix 3848: Unit bays not accepting Aeros for transport, with unit test
MegaMek · Mar 13, 2024 · 0b01443 · 0b01443
2 parents 3ea1787 + 4cf0d4b
commit 0b01443
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 7 deletions.
diff --git a/MekHQ/src/mekhq/campaign/unit/Unit.java b/MekHQ/src/mekhq/campaign/unit/Unit.java
@@ -1283,8 +1283,9 @@ public double getCorrectBayCapacity(int unitType, double unitWeight) {
         switch (unitType) {
             case UnitType.MEK:
                 return getCurrentMechCapacity();
-            case UnitType.AERO:
+            case UnitType.AEROSPACEFIGHTER:
             case UnitType.CONV_FIGHTER:
+            case UnitType.AERO:
                 // Return a small craft slot if no ASF slots exist
                 if (getCurrentASFCapacity() > 0) {
                     return getCurrentASFCapacity();
@@ -1347,8 +1348,9 @@ public void updateBayCapacity(int unitType, double unitWeight, boolean addUnit,
             case UnitType.MEK:
                 setMechCapacity(Math.min((getCurrentMechCapacity() + amount), getMechCapacity()));
                 break;
-            case UnitType.AERO:
+            case UnitType.AEROSPACEFIGHTER:
             case UnitType.CONV_FIGHTER:
+            case UnitType.AERO:
                 // Use the assigned bay number to determine if we need to update ASF or Small Craft capacity
                 Bay aeroBay = getEntity().getBayById(bayNumber);
                 if (aeroBay != null) {
@@ -3420,7 +3422,7 @@ public Color determineForegroundColor(String type) {
             return UIManager.getColor(type + ".Foreground");
         }
     }
-    
+
     public Color determineBackgroundColor(String type) {
         if (isDeployed()) {
             return MekHQ.getMHQOptions().getDeployedBackground();
@@ -3446,7 +3448,7 @@ public Color determineBackgroundColor(String type) {
             return UIManager.getColor(type + ".Background");
         }
     }
-    
+
     /**
      * Determines which crew member is considered the unit commander. For solo-piloted units there is
      * only one option, but units with multiple crew (vehicles, aerospace vessels, infantry) use the following

diff --git a/MekHQ/unittests/mekhq/campaign/unit/UnitTransportTest.java b/MekHQ/unittests/mekhq/campaign/unit/UnitTransportTest.java
@@ -18,12 +18,12 @@
  */
 package mekhq.campaign.unit;
 
-import megamek.common.Aero;
-import megamek.common.Entity;
-import megamek.common.Mech;
+import megamek.common.*;
+import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 
 import java.util.UUID;
+import java.util.Vector;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -32,6 +32,12 @@
 import static org.mockito.Mockito.*;
 
 public class UnitTransportTest {
+
+    @BeforeAll
+    static void before() {
+        EquipmentType.initializeTypes();
+    }
+
     @Test
     public void basicTransportedUnits() {
         Unit transport = new Unit();
@@ -137,6 +143,32 @@ public void isCarryingAeroAndGround() {
         assertFalse(transport.isCarryingGround());
     }
 
+    @Test
+    public void testUnitTypeForAerosMatchesAeroBayType() {
+        // Create a fake entity to back the real transport Unit
+        Dropship mockVengeance = mock(Dropship.class);
+        Unit transport = new Unit();
+        ASFBay mockASFBay = mock(ASFBay.class);
+        when(mockASFBay.getCapacity()).thenReturn(100.0);
+        transport.setEntity(mockVengeance);
+
+        // Initialize bays
+        Vector<Bay> bays = new Vector<>();
+        bays.add(mockASFBay);
+        when(mockVengeance.getTransportBays()).thenReturn(bays);
+        transport.initializeBaySpace();
+
+        // Add an aero unit
+        Entity aero = new AeroSpaceFighter();
+        Unit mockAeroUnit = mock(Unit.class);
+        when(mockAeroUnit.getId()).thenReturn(UUID.randomUUID());
+        when(mockAeroUnit.getEntity()).thenReturn(aero);
+
+        // Verify the AeroSpaceFighter is recognized as a valid ASFBay occupant
+        double remainingCap = transport.getCorrectBayCapacity(aero.getUnitType(), 50);
+        assertEquals(100.0, remainingCap);
+    }
+
     @Test
     public void unloadFromTransportShipDoesNothingIfNotLoaded() {
         Unit transport = spy(new Unit());