Added plague damage support

Author: Bytekeeper

src/main/java/org/bk/ass/Agent.java

@@ -46,6 +46,7 @@ public class Agent {
   int stopFrames;
   int remainingStimFrames;
   boolean canStim;
+  int plagueDamagePerFrameShifted;
 
   // Is Zerg and not an Egg/Larva
   boolean regeneratesHealth;
@@ -337,6 +338,11 @@ public Agent setHpConstructionRate(int buildTime) {
     return this;
   }
 
+  public Agent setPlagueDamage(int plagueDamage) {
+    this.plagueDamagePerFrameShifted = (plagueDamage << 8) / 76;
+    return this;
+  }
+
   public enum TargetingPriority {
     LOW,
     MEDIUM,

src/main/java/org/bk/ass/BWAPI4JAgentFactory.java

@@ -230,7 +230,6 @@ public Agent of(PlayerUnit unit) {
             hasEnergyUpgrade(unitType, player));
     if (map != null && !unit.isFlying()) {
       agent.setElevationLevel(map.getGroundHeight(unit.getTilePosition()));
-//      agent.setProtectedByDarkSwarm(unit.isUnderDarkSwarm()); Not supported by BWAPI4J currently
     }
     if (unitType == UnitType.Terran_Marine || unitType == UnitType.Terran_Firebat) {
       agent.setCanStim(player.hasResearched(TechType.Stim_Packs));

src/main/java/org/bk/ass/BWMirrorAgentFactory.java

@@ -209,7 +209,8 @@ public Agent of(
         .setUserObject(unit)
         .setBurrowed(unit.isBurrowed())
         .setStasised(unit.isStasised())
-        .setLockeddown(unit.isLockedDown());
+        .setLockeddown(unit.isLockedDown())
+        .setPlagueDamage(unit.isPlagued() ? WeaponType.Plague.damageAmount() : 0);
   }
 
   public Agent of(Unit unit) {

src/main/java/org/bk/ass/Simulator.java

@@ -199,6 +199,15 @@ private void updateStats(UnorderedCollection<Agent> agents) {
       if (agent.cooldown > 0) {
         agent.cooldown--;
       }
+      if (agent.shieldsShifted < agent.maxShieldsShifted) {
+        agent.shieldsShifted += 7;
+        if (agent.shieldsShifted > agent.maxShieldsShifted) {
+          agent.shieldsShifted = agent.maxShieldsShifted;
+        }
+      }
+      if (agent.plagueDamagePerFrameShifted < agent.healthShifted) {
+        agent.healthShifted -= agent.plagueDamagePerFrameShifted;
+      }
       if (agent.remainingStimFrames > 0) {
         agent.remainingStimFrames--;
       }
@@ -208,12 +217,6 @@ private void updateStats(UnorderedCollection<Agent> agents) {
           agent.healthShifted = agent.maxHealthShifted;
         }
       }
-      if (agent.shieldsShifted < agent.maxShieldsShifted) {
-        agent.shieldsShifted += 7;
-        if (agent.shieldsShifted > agent.maxShieldsShifted) {
-          agent.shieldsShifted = agent.maxShieldsShifted;
-        }
-      }
       agent.energyShifted += 8;
       if (agent.energyShifted > agent.maxEnergyShifted) {
         agent.energyShifted = agent.maxEnergyShifted;

src/test/java/org/bk/ass/SimulatorTest.java

@@ -5,6 +5,7 @@
 import org.junit.jupiter.api.Test;
 import org.openbw.bwapi4j.test.BWDataProvider;
 import org.openbw.bwapi4j.type.UnitType;
+import org.openbw.bwapi4j.type.WeaponType;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.Assertions.assertThrows;
@@ -771,4 +772,45 @@ void darkSwarmZerglingVsMarines() {
     assertThat(simulator.getAgentsA()).isEmpty();
     assertThat(simulator.getAgentsB()).isNotEmpty();
   }
+
+  @Test
+  void shouldDealPlagueDamage() {
+    // GIVEN
+    simulator.addAgentA(
+        factory.of(UnitType.Terran_Marine).setPlagueDamage(WeaponType.Plague.damageAmount()));
+
+    // Dummy unit
+    simulator.addAgentB(factory.of(UnitType.Protoss_Carrier));
+
+    // WHEN
+    simulator.simulate(1);
+
+    // THEN
+    assertThat(simulator.getAgentsA())
+        .first()
+        .extracting(a -> a.healthShifted)
+        .isEqualTo(9230); // ~4 damage taken
+  }
+
+  @Test
+  void shouldNotDieFromPlague() {
+    // GIVEN
+    simulator.addAgentA(
+        factory
+            .of(UnitType.Terran_Marine)
+            .setPlagueDamage(WeaponType.Plague.damageAmount())
+            .setHealth(2));
+
+    // Dummy unit
+    simulator.addAgentB(factory.of(UnitType.Protoss_Carrier));
+
+    // WHEN
+    simulator.simulate(1);
+
+    // THEN
+    assertThat(simulator.getAgentsA())
+        .first()
+        .extracting(a -> a.healthShifted)
+        .isEqualTo(512); // No damage taken
+  }
 }