Fix colony teams assignment behaviour

src/main/java/com/minecolonies/api/entity/citizen/AbstractEntityCitizen.java

@@ -46,8 +46,10 @@
 import net.minecraft.world.item.ShieldItem;
 import net.minecraft.world.level.Level;
 import net.minecraft.world.phys.Vec3;
+import net.minecraft.world.scores.PlayerTeam;
 import net.minecraftforge.items.IItemHandler;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.List;
 
@@ -233,6 +235,18 @@ public boolean isNoAi()
         return false;
     }
 
+    @Override
+    @Nullable
+    protected PlayerTeam getAssignedTeam()
+    {
+        final ICitizenColonyHandler citizenColonyHandler = getCitizenColonyHandler();
+        if (citizenColonyHandler == null || citizenColonyHandler.getColony() == null)
+        {
+            return null;
+        }
+        return citizenColonyHandler.getColony().getTeam();
+    }
+
     /**
      * Sets the textures of all citizens and distinguishes between male and female.
      */

src/main/java/com/minecolonies/api/entity/mobs/AbstractEntityRaiderMob.java

@@ -52,6 +52,7 @@
 import static com.minecolonies.api.util.constant.ColonyManagerConstants.NO_COLONY_ID;
 import static com.minecolonies.api.util.constant.NbtTagConstants.*;
 import static com.minecolonies.api.util.constant.RaiderConstants.*;
+import static com.minecolonies.core.util.TeamUtils.checkOrCreateTeam;
 
 /**
  * Abstract for all raider entities.
@@ -521,23 +522,23 @@ private void onEnterChunk(final ChunkPos newChunkPos)
         }
     }
 
-    @org.jetbrains.annotations.Nullable
+    @Nullable
     @Override
     public SpawnGroupData finalizeSpawn(
       final ServerLevelAccessor worldIn,
       final DifficultyInstance difficultyIn,
       final MobSpawnType reason,
-      @org.jetbrains.annotations.Nullable final SpawnGroupData spawnDataIn,
-      @org.jetbrains.annotations.Nullable final CompoundTag dataTag)
+      @Nullable final SpawnGroupData spawnDataIn,
+      @Nullable final CompoundTag dataTag)
     {
         RaiderMobUtils.setEquipment(this);
         return super.finalizeSpawn(worldIn, difficultyIn, reason, spawnDataIn, dataTag);
     }
 
     @Override
-    public void remove(RemovalReason reason)
+    public void remove(@NotNull final RemovalReason reason)
     {
-        if (!level().isClientSide && colony != null && eventID > 0)
+        if (!level.isClientSide && colony != null && eventID > 0)
         {
             colony.getEventManager().unregisterEntity(this, eventID);
         }
@@ -732,38 +733,15 @@ public void initStatsFor(final double baseHealth, final double difficulty, final
             this.setEnvDamageImmunity(true);
         }
 
-        if (difficulty >= TEAM_DIFFICULTY)
-        {
-            level().getScoreboard().addPlayerToTeam(getScoreboardName(), checkOrCreateTeam());
-        }
-
         this.getAttribute(Attributes.MAX_HEALTH).setBaseValue(baseHealth);
         this.setHealth(this.getMaxHealth());
     }
 
-    /**
-     * Creates or gets the scoreboard team
-     *
-     * @return Scoreboard team
-     */
-    private PlayerTeam checkOrCreateTeam()
-    {
-        if (this.level().getScoreboard().getPlayerTeam(getTeamName()) == null)
-        {
-            this.level().getScoreboard().addPlayerTeam(getTeamName());
-            this.level().getScoreboard().getPlayerTeam(getTeamName()).setAllowFriendlyFire(false);
-        }
-        return this.level().getScoreboard().getPlayerTeam(getTeamName());
-    }
-
-    /**
-     * Gets the scoreboard team name
-     *
-     * @return
-     */
-    protected String getTeamName()
+    @Override
+    @Nullable
+    protected PlayerTeam getAssignedTeam()
     {
-        return RAID_TEAM;
+        return checkOrCreateTeam(level, RAID_TEAM);
     }
 
     /**

src/main/java/com/minecolonies/api/entity/other/AbstractFastMinecoloniesEntity.java

@@ -14,12 +14,15 @@
 import net.minecraft.world.level.Level;
 import net.minecraft.world.level.entity.EntityTypeTest;
 import net.minecraft.world.phys.Vec3;
+import net.minecraft.world.scores.PlayerTeam;
+import net.minecraft.world.scores.Team;
 import net.minecraftforge.common.util.ITeleporter;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 
 /**
  * Special abstract minecolonies mob that overrides laggy vanilla behaviour.
@@ -311,6 +314,74 @@ public void updateSwimAmount()
 
     }
 
+    /**
+     * Get the team this entity is assigned to.
+     *
+     * @return the team instance.
+     */
+    @Nullable
+    protected abstract PlayerTeam getAssignedTeam();
+
+    @Override
+    @Nullable
+    public final Team getTeam()
+    {
+        final PlayerTeam assignedTeam = getAssignedTeam();
+        registerToTeamInternal(assignedTeam);
+        return assignedTeam;
+    }
+
+    /**
+     * Register this entity to its own assigned team.
+     */
+    public void registerToTeam()
+    {
+        registerToTeamInternal(getAssignedTeam());
+    }
+
+    /**
+     * Internal method for team registration.
+     *
+     * @param team the team to register to.
+     */
+    private void registerToTeamInternal(@Nullable final PlayerTeam team)
+    {
+        if (team != null && !isInTeam(team))
+        {
+            level.getScoreboard().addPlayerToTeam(getScoreboardName(), team);
+        }
+    }
+
+    /**
+     * Remove the entity from its own assigned team.
+     */
+    public void removeFromTeam()
+    {
+        final PlayerTeam team = getAssignedTeam();
+        if (team != null && isInTeam(team))
+        {
+            level.getScoreboard().removePlayerFromTeam(getScoreboardName(), team);
+        }
+    }
+
+    /**
+     * Check if the current entity is assigned to the provided team.
+     *
+     * @param team the input team.
+     * @return true if so.
+     */
+    private boolean isInTeam(@NotNull final PlayerTeam team)
+    {
+        return Objects.equals(level.getScoreboard().getPlayersTeam(getScoreboardName()), team);
+    }
+
+    @Override
+    public void remove(@NotNull final RemovalReason reason)
+    {
+        super.remove(reason);
+        removeFromTeam();
+    }
+
     /**
      * Static Byte values to avoid frequent autoboxing
      */

src/main/java/com/minecolonies/core/colony/Colony.java

@@ -84,6 +84,7 @@
 import static com.minecolonies.api.util.constant.NbtTagConstants.*;
 import static com.minecolonies.api.util.constant.TranslationConstants.*;
 import static com.minecolonies.core.MineColonies.getConfig;
+import static com.minecolonies.core.util.TeamUtils.checkOrCreateTeam;
 
 /**
  * This class describes a colony and contains all the data and methods for manipulating a Colony.
@@ -366,7 +367,7 @@ protected Colony(final int id, @Nullable final Level world)
         {
             this.dimensionId = world.dimension();
             onWorldLoad(world);
-            checkOrCreateTeam();
+            checkOrCreateTeam(world, getTeamName());
         }
         this.permissions = new Permissions(this);
         researchManager = new ResearchManager(this);
@@ -631,23 +632,11 @@ public void updateAttackingPlayers()
     }
 
     @Override
+    @Nullable
     public PlayerTeam getTeam()
     {
         // This getter will create the team if it doesn't exist. Could do something different though in the future.
-        return checkOrCreateTeam();
-    }
-
-    /**
-     * Check or create the team.
-     */
-    private PlayerTeam checkOrCreateTeam()
-    {
-        if (this.world.getScoreboard().getPlayerTeam(getTeamName()) == null)
-        {
-            this.world.getScoreboard().addPlayerTeam(getTeamName());
-            this.world.getScoreboard().getPlayerTeam(getTeamName()).setAllowFriendlyFire(false);
-        }
-        return this.world.getScoreboard().getPlayerTeam(getTeamName());
+        return checkOrCreateTeam(world, getTeamName());
     }
 
     /**
@@ -659,10 +648,13 @@ public void setColonyColor(final ChatFormatting colonyColor)
     {
         if (this.world != null)
         {
-            checkOrCreateTeam();
             this.colonyTeamColor = colonyColor;
-            this.world.getScoreboard().getPlayerTeam(getTeamName()).setColor(colonyColor);
-            this.world.getScoreboard().getPlayerTeam(getTeamName()).setPlayerPrefix(Component.literal(colonyColor.toString()));
+            final PlayerTeam team = getTeam();
+            if (team != null)
+            {
+                team.setColor(colonyColor);
+                team.setPlayerPrefix(Component.literal(colonyColor.toString()));
+            }
         }
         this.markDirty();
     }