diff --git a/fp/battle_modifier.py b/fp/battle_modifier.py index 1adf8a38..57083776 100644 --- a/fp/battle_modifier.py +++ b/fp/battle_modifier.py @@ -992,6 +992,15 @@ def activate(battle, split_msg): logger.info("Setting {}'s item to {}".format(pkmn.name, item)) pkmn.item = item + if split_msg[3].lower().startswith("move: "): + move_name = normalize_name(split_msg[3].split(":")[-1].strip()) + if ( + move_name in all_move_json + and all_move_json[move_name].get("volatileStatus") == "partiallytrapped" + ): + logger.info("{} was partially trapped by {}".format(pkmn.name, move_name)) + pkmn.volatile_statuses.append("partiallytrapped") + def anim(battle, split_msg): if is_opponent(battle, split_msg): @@ -1139,6 +1148,8 @@ def end_volatile_status(battle, split_msg): if vs.startswith(volatile_status): logger.info("Removing {} from {}".format(vs, pkmn.name)) pkmn.volatile_statuses.remove(vs) + elif len(split_msg) >= 5 and "partiallytrapped" in split_msg[4]: + remove_volatile(pkmn, "partiallytrapped") elif volatile_status not in pkmn.volatile_statuses: logger.warning( "Pokemon '{}' does not have the volatile status '{}'".format( diff --git a/tests/test_battle_modifiers.py b/tests/test_battle_modifiers.py index a3cc720f..bdf1c22a 100644 --- a/tests/test_battle_modifiers.py +++ b/tests/test_battle_modifiers.py @@ -1025,6 +1025,42 @@ def setUp(self): self.battle.opponent.active = self.opponent_active self.battle.user.active = self.user_active + def test_activating_partially_trapped_whirlpool(self): + split_msg = [ + "", + "-activate", + "p2a: Caterpie", + "move: Whirlpool", + "[of] p1a: Luvdisc", + ] + activate(self.battle, split_msg) + self.assertIn("partiallytrapped", self.battle.opponent.active.volatile_statuses) + + def test_activating_partially_trapped_magmastorm(self): + split_msg = [ + "", + "-activate", + "p2a: Caterpie", + "move: Magma Storm", + "[of] p1a: Luvdisc", + ] + activate(self.battle, split_msg) + self.assertIn("partiallytrapped", self.battle.opponent.active.volatile_statuses) + + def test_does_not_activate_partiallytrapped_when_not_a_partiallytrapping_move(self): + # this isn't something that would cause an `-activate`, but just to make sure the logic is correct + split_msg = [ + "", + "-activate", + "p2a: Caterpie", + "move: Tackle", + "[of] p1a: Luvdisc", + ] + activate(self.battle, split_msg) + self.assertNotIn( + "partiallytrapped", self.battle.opponent.active.volatile_statuses + ) + def test_sets_item_when_poltergeist_activates(self): split_msg = [ "", @@ -2470,6 +2506,27 @@ def setUp(self): self.user_active = Pokemon("weedle", 100) self.battle.user.active = self.user_active + def test_removes_partiallytrapped(self): + self.battle.opponent.active.volatile_statuses = ["partiallytrapped"] + split_msg = ["", "-end", "p2a: Caterpie", "whirlpool", "[partiallytrapped]"] + end_volatile_status(self.battle, split_msg) + + self.assertEqual([], self.battle.opponent.active.volatile_statuses) + + def test_removes_partiallytrapped_silent(self): + self.battle.opponent.active.volatile_statuses = ["partiallytrapped"] + split_msg = [ + "", + "-end", + "p2a: Caterpie", + "whirlpool", + "[partiallytrapped]", + "[silent]", + ] + end_volatile_status(self.battle, split_msg) + + self.assertEqual([], self.battle.opponent.active.volatile_statuses) + def test_removes_volatile_status_from_opponent(self): self.battle.opponent.active.volatile_statuses = ["encore"] split_msg = ["", "-end", "p2a: Caterpie", "Encore"]