A bit of cleanup

src/main/java/com/minelittlepony/unicopia/entity/behaviour/BeeBehaviour.java

@@ -1,7 +1,8 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
+
 import net.minecraft.client.MinecraftClient;
 import net.minecraft.client.sound.AggressiveBeeSoundInstance;
 import net.minecraft.client.sound.PassiveBeeSoundInstance;
@@ -20,7 +21,7 @@ public class BeeBehaviour extends EntityBehaviour<BeeEntity> {
     }
 
     @Override
-    public void update(Caster<?> source, BeeEntity entity, Spell spell) {
+    public void update(Caster<?> source, BeeEntity entity, DisguiseSpell spell) {
 
         if (source.getOwner().isSneaking()) {
             entity.setAngerTime(10);

src/main/java/com/minelittlepony/unicopia/entity/behaviour/ChickenBehaviour.java

@@ -1,7 +1,7 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
 import com.minelittlepony.unicopia.entity.player.Pony;
 
 import net.minecraft.entity.Entity;
@@ -21,7 +21,7 @@ public class ChickenBehaviour extends EntityBehaviour<ChickenEntity> {
     }
 
     @Override
-    public void update(Caster<?> source, ChickenEntity entity, Spell spell) {
+    public void update(Caster<?> source, ChickenEntity entity, DisguiseSpell spell) {
         entity.eggLayTime = Integer.MAX_VALUE;
 
         if (source instanceof Pony) {

src/main/java/com/minelittlepony/unicopia/entity/behaviour/CreeperBehaviour.java

@@ -1,13 +1,13 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
 
 import net.minecraft.entity.mob.CreeperEntity;
 
 public class CreeperBehaviour extends EntityBehaviour<CreeperEntity> {
     @Override
-    public void update(Caster<?> source, CreeperEntity entity, Spell spell) {
+    public void update(Caster<?> source, CreeperEntity entity, DisguiseSpell spell) {
         if (isSneakingOnGround(source)) {
             entity.setFuseSpeed(1);
         } else {

src/main/java/com/minelittlepony/unicopia/entity/behaviour/EndermanBehaviour.java

@@ -1,7 +1,7 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
 
 import net.minecraft.entity.mob.EndermanEntity;
 import net.minecraft.item.BlockItem;
@@ -10,7 +10,7 @@ import net.minecraft.util.Hand;
 
 public class EndermanBehaviour extends EntityBehaviour<EndermanEntity> {
     @Override
-    public void update(Caster<?> source, EndermanEntity entity, Spell spell) {
+    public void update(Caster<?> source, EndermanEntity entity, DisguiseSpell spell) {
         if (source.getOwner().isSneaking() || source.getOwner().isSprinting()) {
             entity.setTarget(entity);
         } else {

src/main/java/com/minelittlepony/unicopia/entity/behaviour/EntityBehaviour.java

@@ -6,8 +6,9 @@ import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
 import com.minelittlepony.unicopia.entity.ItemWielder;
+import com.minelittlepony.unicopia.entity.player.Pony;
 import com.minelittlepony.unicopia.util.Registries;
 
 import net.minecraft.entity.Entity;
@@ -35,7 +36,13 @@ public class EntityBehaviour<T extends Entity> {
     private static final EntityBehaviour<Entity> DEFAULT = new EntityBehaviour<>();
     private static final Registry<EntityBehaviour<?>> REGISTRY = Registries.createSimple(new Identifier("unicopia", "entity_behaviour"));
 
-    public void update(Caster<?> source, T entity, Spell spell) {
+    public void update(Caster<?> source, T entity, DisguiseSpell spell) {
+        if (source instanceof Pony) {
+            update((Pony)source, entity, spell);
+        }
+    }
+
+    protected void update(Pony pony, T entity, DisguiseSpell spell) {
 
     }
 
@@ -224,7 +231,7 @@ public class EntityBehaviour<T extends Entity> {
         register(FallingBlockBehaviour::new, EntityType.FALLING_BLOCK);
         register(MobBehaviour::new, EntityType.RAVAGER, EntityType.IRON_GOLEM);
         register(RabbitBehaviour::new, EntityType.RABBIT);
-        register(VillagerBehaviour::new, EntityType.VILLAGER, EntityType.WANDERING_TRADER);
+        register(TraderBehaviour::new, EntityType.VILLAGER, EntityType.WANDERING_TRADER);
         register(SheepBehaviour::new, EntityType.SHEEP);
         register(BeeBehaviour::new, EntityType.BEE);
         register(EndermanBehaviour::new, EntityType.ENDERMAN);

src/main/java/com/minelittlepony/unicopia/entity/behaviour/FallingBlockBehaviour.java

@@ -4,7 +4,6 @@ import java.util.List;
 import java.util.Optional;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
 import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
 import com.minelittlepony.unicopia.entity.player.Pony;
 import com.minelittlepony.unicopia.mixin.MixinBlockEntity;
@@ -68,7 +67,7 @@ public class FallingBlockBehaviour extends EntityBehaviour<FallingBlockEntity> {
     }
 
     @Override
-    public void update(Caster<?> source, FallingBlockEntity entity, Spell spell) {
+    public void update(Caster<?> source, FallingBlockEntity entity, DisguiseSpell spell) {
 
         BlockState state = entity.getBlockState();
         if (state.contains(Properties.WATERLOGGED)) {
@@ -76,12 +75,12 @@ public class FallingBlockBehaviour extends EntityBehaviour<FallingBlockEntity> {
 
             if (state.get(Properties.WATERLOGGED) != logged) {
                 entity = new FallingBlockEntity(entity.world, entity.getX(), entity.getY(), entity.getZ(), state.with(Properties.WATERLOGGED, logged));
-                ((DisguiseSpell)spell).getDisguise().setAppearance(entity);
+                spell.getDisguise().setAppearance(entity);
                 return;
             }
         }
 
-        Disguise disguise = ((DisguiseSpell)spell).getDisguise();
+        Disguise disguise = spell.getDisguise();
         List<Entity> attachments = disguise.getAttachments();
         if (attachments.size() > 0) {
             copyBaseAttributes(source.getOwner(), attachments.get(0), UP);

src/main/java/com/minelittlepony/unicopia/entity/behaviour/MinecartBehaviour.java

@@ -1,7 +1,7 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
 
 import net.minecraft.client.MinecraftClient;
 import net.minecraft.client.sound.MovingMinecartSoundInstance;
@@ -20,7 +20,7 @@ public class MinecartBehaviour extends EntityBehaviour<AbstractMinecartEntity> {
     }
 
     @Override
-    public void update(Caster<?> source, AbstractMinecartEntity entity, Spell spell) {
+    public void update(Caster<?> source, AbstractMinecartEntity entity, DisguiseSpell spell) {
         entity.yaw -= 90;
         entity.prevYaw -= 90;
 

src/main/java/com/minelittlepony/unicopia/entity/behaviour/MobBehaviour.java

@@ -1,7 +1,6 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
-import com.minelittlepony.unicopia.ability.magic.Caster;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
-import com.minelittlepony.unicopia.ability.magic.Spell;
 import com.minelittlepony.unicopia.entity.player.Pony;
 import com.minelittlepony.unicopia.util.RayTraceHelper;
 
@@ -10,14 +9,10 @@ import net.minecraft.entity.mob.MobEntity;
 
 public class MobBehaviour extends EntityBehaviour<MobEntity> {
     @Override
-    public void update(Caster<?> source, MobEntity entity, Spell spell) {
+    public void update(Pony player, MobEntity entity, DisguiseSpell spell) {
 
-        if (source instanceof Pony) {
+        if (player.sneakingChanged() && isSneakingOnGround(player)) {
-            Pony player = (Pony)source;
+            entity.tryAttack(RayTraceHelper.findEntity(player.getEntity(), 6, 1, e -> e instanceof LivingEntity).orElse(entity));
-
-            if (player.sneakingChanged() && isSneakingOnGround(source)) {
-                entity.tryAttack(RayTraceHelper.findEntity(source.getEntity(), 6, 1, e -> e instanceof LivingEntity).orElse(entity));
-            }
         }
     }
 }

src/main/java/com/minelittlepony/unicopia/entity/behaviour/RabbitBehaviour.java

@@ -1,7 +1,6 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
-import com.minelittlepony.unicopia.ability.magic.Caster;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
-import com.minelittlepony.unicopia.ability.magic.Spell;
 import com.minelittlepony.unicopia.entity.player.Pony;
 
 import net.minecraft.entity.Entity;
@@ -9,19 +8,15 @@ import net.minecraft.entity.passive.RabbitEntity;
 
 public class RabbitBehaviour extends EntityBehaviour<RabbitEntity> {
     @Override
-    public void update(Caster<?> source, RabbitEntity entity, Spell spell) {
+    public void update(Pony player, RabbitEntity entity, DisguiseSpell spell) {
 
-        if (source instanceof Pony) {
+        if (player.getEntity().isOnGround()) {
-            Pony player = (Pony)source;
+            if (Entity.squaredHorizontalLength(player.getEntity().getVelocity()) > 0.01) {
-
+                player.getOwner().jump();
-            if (player.getEntity().isOnGround()) {
-                if (Entity.squaredHorizontalLength(player.getEntity().getVelocity()) > 0.01) {
-                    player.getOwner().jump();
-                    entity.startJump();
-                }
-            } else if (player.landedChanged()) {
                 entity.startJump();
             }
+        } else if (player.landedChanged()) {
+            entity.startJump();
         }
     }
 }

src/main/java/com/minelittlepony/unicopia/entity/behaviour/RangedAttackBehaviour.java

@@ -1,8 +1,7 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
 import java.util.function.BiFunction;
-import com.minelittlepony.unicopia.ability.magic.Caster;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
-import com.minelittlepony.unicopia.ability.magic.Spell;
 import com.minelittlepony.unicopia.entity.player.Pony;
 
 import net.minecraft.entity.Entity;
@@ -24,28 +23,24 @@ public class RangedAttackBehaviour<T extends Entity & RangedAttackMob> extends E
     }
 
     @Override
-    public void update(Caster<?> source, T entity, Spell spell) {
+    public void update(Pony player, T entity, DisguiseSpell spell) {
 
-        if (source instanceof Pony) {
+        if (player.sneakingChanged() && isSneakingOnGround(player)) {
-            Pony player = (Pony)source;
 
-            if (player.sneakingChanged() && isSneakingOnGround(player)) {
+            ProjectileEntity spit = projectileSupplier.apply(entity.world, entity);
 
-                ProjectileEntity spit = projectileSupplier.apply(entity.world, entity);
+            Vec3d rot = player.getEntity().getRotationVec(1);
 
-                Vec3d rot = source.getEntity().getRotationVec(1);
+            spit.setVelocity(rot.getX(), rot.getY(), rot.getZ(), 1.5F, 3);
+            spit.setOwner(player.getOwner());
 
-                spit.setVelocity(rot.getX(), rot.getY(), rot.getZ(), 1.5F, 3);
+            if (!entity.isSilent()) {
-                spit.setOwner(source.getOwner());
+                entity.world.playSound(null, entity.getX(), entity.getY(), entity.getZ(),
-
+                        sound, entity.getSoundCategory(), 1,
-                if (!entity.isSilent()) {
+                        1 + (entity.world.random.nextFloat() - entity.world.random.nextFloat()) * 0.2F);
-                    entity.world.playSound(null, entity.getX(), entity.getY(), entity.getZ(),
-                            sound, entity.getSoundCategory(), 1,
-                            1 + (entity.world.random.nextFloat() - entity.world.random.nextFloat()) * 0.2F);
-                }
-
-                entity.world.spawnEntity(spit);
             }
+
+            entity.world.spawnEntity(spit);
         }
     }
 }

src/main/java/com/minelittlepony/unicopia/entity/behaviour/SheepBehaviour.java

@@ -2,8 +2,7 @@ package com.minelittlepony.unicopia.entity.behaviour;
 
 import java.util.Random;
 
-import com.minelittlepony.unicopia.ability.magic.Caster;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
-import com.minelittlepony.unicopia.ability.magic.Spell;
 import com.minelittlepony.unicopia.entity.player.Pony;
 import com.minelittlepony.unicopia.mixin.MixinSheepEntity;
 import com.minelittlepony.unicopia.util.WorldEvent;
@@ -20,57 +19,54 @@ import net.minecraft.util.math.BlockPos;
 
 public class SheepBehaviour extends EntityBehaviour<SheepEntity> {
     @Override
-    public void update(Caster<?> source, SheepEntity entity, Spell spell) {
+    public void update(Pony player, SheepEntity entity, DisguiseSpell spell) {
-        if (source instanceof Pony) {
-            Pony player = (Pony)source;
 
-            if (player.sneakingChanged()) {
+        if (player.sneakingChanged()) {
 
-                BlockPos pos = entity.getBlockPos().down();
+            BlockPos pos = entity.getBlockPos().down();
-                BlockState state = entity.world.getBlockState(pos);
+            BlockState state = entity.world.getBlockState(pos);
-                boolean grass = state.isOf(Blocks.GRASS_BLOCK);
+            boolean grass = state.isOf(Blocks.GRASS_BLOCK);
 
-                if (player.getOwner().isSneaking()) {
+            if (player.getOwner().isSneaking()) {
-                    if (grass && entity.world.isClient && entity.isSheared()) {
+                if (grass && entity.world.isClient && entity.isSheared()) {
-                        entity.handleStatus((byte)10);
+                    entity.handleStatus((byte)10);
-                    }
-                } else {
-                    if (entity.isSheared() && grass) {
-                        WorldEvent.play(WorldEvent.DESTROY_BLOCK, entity.world, pos, state);
-                        entity.world.setBlockState(pos, Blocks.DIRT.getDefaultState(), 2);
-
-                        entity.onEatingGrass();
-                    } else if (!entity.isSheared()) {
-                        ItemStack dropType = new ItemStack(MixinSheepEntity.getDrops().get(entity.getColor()).asItem());
-
-                        player.getOwner().playSound(SoundEvents.ENTITY_SHEEP_SHEAR, SoundCategory.PLAYERS, 1, 1);
-                        entity.setSheared(true);
-
-                        Random rng = entity.world.random;
-                        PlayerInventory inv = player.getOwner().inventory;
-
-                        int dropAmount = rng.nextInt(3);
-                        int slot;
-
-                        do {
-                            slot = inv.method_7371(dropType);
-                            if (slot < 0) {
-                                break;
-                            }
-                            inv.removeStack(slot, 1);
-                            ItemEntity itemEntity = entity.dropItem(dropType.getItem(), 1);
-                            if (itemEntity != null) {
-                               itemEntity.setVelocity(itemEntity.getVelocity().add(
-                                       (rng.nextFloat() - rng.nextFloat()) * 0.1F,
-                                       rng.nextFloat() * 0.05F,
-                                       (rng.nextFloat() - rng.nextFloat()) * 0.1F
-                               ));
-                               itemEntity.setPickupDelay(40);
-                            }
-                        } while (dropAmount-- > 0);
-                    }
-                    spell.setDirty(true);
                 }
+            } else {
+                if (entity.isSheared() && grass) {
+                    WorldEvent.play(WorldEvent.DESTROY_BLOCK, entity.world, pos, state);
+                    entity.world.setBlockState(pos, Blocks.DIRT.getDefaultState(), 2);
+
+                    entity.onEatingGrass();
+                } else if (!entity.isSheared()) {
+                    ItemStack dropType = new ItemStack(MixinSheepEntity.getDrops().get(entity.getColor()).asItem());
+
+                    player.getOwner().playSound(SoundEvents.ENTITY_SHEEP_SHEAR, SoundCategory.PLAYERS, 1, 1);
+                    entity.setSheared(true);
+
+                    Random rng = entity.world.random;
+                    PlayerInventory inv = player.getOwner().inventory;
+
+                    int dropAmount = rng.nextInt(3);
+                    int slot;
+
+                    do {
+                        slot = inv.method_7371(dropType);
+                        if (slot < 0) {
+                            break;
+                        }
+                        inv.removeStack(slot, 1);
+                        ItemEntity itemEntity = entity.dropItem(dropType.getItem(), 1);
+                        if (itemEntity != null) {
+                           itemEntity.setVelocity(itemEntity.getVelocity().add(
+                                   (rng.nextFloat() - rng.nextFloat()) * 0.1F,
+                                   rng.nextFloat() * 0.05F,
+                                   (rng.nextFloat() - rng.nextFloat()) * 0.1F
+                           ));
+                           itemEntity.setPickupDelay(40);
+                        }
+                    } while (dropAmount-- > 0);
+                }
+                spell.setDirty(true);
             }
         }
     }

src/main/java/com/minelittlepony/unicopia/entity/behaviour/ShulkerBehaviour.java

@@ -1,7 +1,7 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
 import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
 import com.minelittlepony.unicopia.entity.player.Pony;
 import com.minelittlepony.unicopia.mixin.MixinShulkerEntity;
 
@@ -16,38 +16,14 @@ import net.minecraft.util.math.Vec3d;
 
 public class ShulkerBehaviour extends EntityBehaviour<ShulkerEntity> {
     @Override
-    public void update(Caster<?> source, ShulkerEntity shulker, Spell spell) {
+    public void update(Caster<?> source, ShulkerEntity shulker, DisguiseSpell spell) {
         shulker.yaw = 0;
         shulker.prevBodyYaw = 0;
         shulker.bodyYaw = 0;
 
         shulker.setAttachedBlock(null);
 
-        if (source instanceof Pony) {
+        super.update(source, shulker, spell);
-            Pony player = (Pony)source;
-
-            float peekAmount = 30;
-
-            double speed = !source.getEntity().isSneaking() ? 0.29 : 0;
-            speed += Math.sqrt(Entity.squaredHorizontalLength(source.getEntity().getVelocity())) * 2;
-
-            peekAmount = (float)MathHelper.clamp(speed, 0, 1);
-            peekAmount = ((Pony)source).getInterpolator().interpolate("peek", peekAmount, 5);
-
-            MixinShulkerEntity mx = (MixinShulkerEntity)shulker;
-
-            mx.setPrevOpenProgress(mx.getOpenProgress());
-            mx.setOpenProgress(peekAmount);
-
-            if (player.sneakingChanged()) {
-                shulker.setPeekAmount((int)(peekAmount / 0.01F));
-            } else if (peekAmount > 0.2 && shulker.getPeekAmount() == 0) {
-                if (shulker.isAlive() && shulker.world.random.nextInt(1000) < shulker.ambientSoundChance++) {
-                    shulker.ambientSoundChance = -shulker.getMinAmbientSoundDelay();
-                    shulker.playSound(SoundEvents.ENTITY_SHULKER_AMBIENT, 1, 1);
-                 }
-            }
-        }
 
         Direction attachmentFace = shulker.getAttachedFace();
         BlockPos pos = shulker.getBlockPos().offset(attachmentFace);
@@ -63,4 +39,29 @@ public class ShulkerBehaviour extends EntityBehaviour<ShulkerEntity> {
             }
         }
     }
+
+    @Override
+    protected void update(Pony player, ShulkerEntity shulker, DisguiseSpell spell) {
+        float peekAmount = 30;
+
+        double speed = !player.getEntity().isSneaking() ? 0.29 : 0;
+        speed += Math.sqrt(Entity.squaredHorizontalLength(player.getEntity().getVelocity())) * 2;
+
+        peekAmount = (float)MathHelper.clamp(speed, 0, 1);
+        peekAmount = player.getInterpolator().interpolate("peek", peekAmount, 5);
+
+        MixinShulkerEntity mx = (MixinShulkerEntity)shulker;
+
+        mx.setPrevOpenProgress(mx.getOpenProgress());
+        mx.setOpenProgress(peekAmount);
+
+        if (player.sneakingChanged()) {
+            shulker.setPeekAmount((int)(peekAmount / 0.01F));
+        } else if (peekAmount > 0.2 && shulker.getPeekAmount() == 0) {
+            if (shulker.isAlive() && shulker.world.random.nextInt(1000) < shulker.ambientSoundChance++) {
+                shulker.ambientSoundChance = -shulker.getMinAmbientSoundDelay();
+                shulker.playSound(SoundEvents.ENTITY_SHULKER_AMBIENT, 1, 1);
+             }
+        }
+    }
 }

src/main/java/com/minelittlepony/unicopia/entity/behaviour/SilverfishBehaviour.java

@@ -1,7 +1,6 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
-import com.minelittlepony.unicopia.ability.magic.Caster;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
-import com.minelittlepony.unicopia.ability.magic.Spell;
 import com.minelittlepony.unicopia.block.state.StateMaps;
 import com.minelittlepony.unicopia.entity.player.Pony;
 import com.minelittlepony.unicopia.util.WorldEvent;
@@ -12,19 +11,15 @@ import net.minecraft.util.math.BlockPos;
 
 public class SilverfishBehaviour extends EntityBehaviour<SilverfishEntity> {
     @Override
-    public void update(Caster<?> source, SilverfishEntity entity, Spell spell) {
+    public void update(Pony player, SilverfishEntity entity, DisguiseSpell spell) {
-        if (source instanceof Pony && !source.isClient()) {
+        if (!player.isClient() && player.sneakingChanged() && player.getOwner().isSneaking()) {
-            Pony player = (Pony)source;
+            BlockPos pos = entity.getBlockPos().down();
+            BlockState state = entity.world.getBlockState(pos);
 
-            if (player.sneakingChanged() && player.getOwner().isSneaking()) {
+            if (StateMaps.SILVERFISH_AFFECTED.canConvert(state)) {
-                BlockPos pos = entity.getBlockPos().down();
-                BlockState state = entity.world.getBlockState(pos);
 
-                if (StateMaps.SILVERFISH_AFFECTED.canConvert(state)) {
+                entity.world.setBlockState(pos, StateMaps.SILVERFISH_AFFECTED.getConverted(state));
-
+                WorldEvent.play(WorldEvent.DESTROY_BLOCK, entity.world, pos, state);
-                    entity.world.setBlockState(pos, StateMaps.SILVERFISH_AFFECTED.getConverted(state));
-                    WorldEvent.play(WorldEvent.DESTROY_BLOCK, entity.world, pos, state);
-                }
             }
         }
     }

src/main/java/com/minelittlepony/unicopia/entity/behaviour/SpellcastingIllagerBehaviour.java

@@ -1,29 +1,23 @@
 package com.minelittlepony.unicopia.entity.behaviour;
 
-import com.minelittlepony.unicopia.ability.magic.Caster;
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
-import com.minelittlepony.unicopia.ability.magic.Spell;
 import com.minelittlepony.unicopia.entity.player.Pony;
 
 import net.minecraft.entity.mob.SpellcastingIllagerEntity;
 
 public class SpellcastingIllagerBehaviour extends EntityBehaviour<SpellcastingIllagerEntity> {
     @Override
-    public void update(Caster<?> source, SpellcastingIllagerEntity entity, Spell s) {
+    public void update(Pony player, SpellcastingIllagerEntity entity, DisguiseSpell s) {
+        if (player.sneakingChanged()) {
+            if (player.getOwner().isSneaking()) {
+                SpellcastingIllagerEntity.Spell[] spells = SpellcastingIllagerEntity.Spell.values();
+                SpellcastingIllagerEntity.Spell spell = spells[entity.world.random.nextInt(spells.length - 1) + 1];
 
-        if (source instanceof Pony) {
+                entity.setSpell(spell);
-            Pony player = (Pony)source;
+                entity.setTarget(entity);
-
+            } else {
-            if (player.sneakingChanged()) {
+                entity.setSpell(SpellcastingIllagerEntity.Spell.NONE);
-                if (player.getOwner().isSneaking()) {
+                entity.setTarget(null);
-                    SpellcastingIllagerEntity.Spell[] spells = SpellcastingIllagerEntity.Spell.values();
-                    SpellcastingIllagerEntity.Spell spell = spells[entity.world.random.nextInt(spells.length - 1) + 1];
-
-                    entity.setSpell(spell);
-                    entity.setTarget(entity);
-                } else {
-                    entity.setSpell(SpellcastingIllagerEntity.Spell.NONE);
-                    entity.setTarget(null);
-                }
             }
         }
     }

src/main/java/com/minelittlepony/unicopia/entity/behaviour/TraderBehaviour.java

@@ -0,0 +1,20 @@
+package com.minelittlepony.unicopia.entity.behaviour;
+
+import com.minelittlepony.unicopia.ability.magic.spell.DisguiseSpell;
+import com.minelittlepony.unicopia.entity.player.Pony;
+
+import net.minecraft.entity.passive.AbstractTraderEntity;
+import net.minecraft.sound.SoundEvents;
+
+public class TraderBehaviour extends EntityBehaviour<AbstractTraderEntity> {
+    @Override
+    public void update(Pony pony, AbstractTraderEntity entity, DisguiseSpell spell) {
+        if (pony.sneakingChanged() && pony.getOwner().isSneaking()) {
+            entity.setHeadRollingTimeLeft(40);
+
+            if (!entity.world.isClient()) {
+               entity.playSound(SoundEvents.ENTITY_VILLAGER_NO, 1, 1);
+            }
+        }
+    }
+}

src/main/java/com/minelittlepony/unicopia/entity/behaviour/VillagerBehaviour.java

@@ -1,26 +0,0 @@
-package com.minelittlepony.unicopia.entity.behaviour;
-
-import com.minelittlepony.unicopia.ability.magic.Caster;
-import com.minelittlepony.unicopia.ability.magic.Spell;
-import com.minelittlepony.unicopia.entity.player.Pony;
-
-import net.minecraft.entity.passive.AbstractTraderEntity;
-import net.minecraft.sound.SoundEvents;
-
-public class VillagerBehaviour extends EntityBehaviour<AbstractTraderEntity> {
-    @Override
-    public void update(Caster<?> source, AbstractTraderEntity entity, Spell spell) {
-
-        if (source instanceof Pony) {
-            Pony player = (Pony)source;
-
-            if (player.sneakingChanged() && player.getOwner().isSneaking()) {
-                entity.setHeadRollingTimeLeft(40);
-
-                if (!entity.world.isClient()) {
-                   entity.playSound(SoundEvents.ENTITY_VILLAGER_NO, 1, 1);
-                }
-            }
-        }
-    }
-}