From 8238bd611caaabe1e0307513b77f52956b1b2238 Mon Sep 17 00:00:00 2001
From: Archy-X <63976867+Archy-X@users.noreply.github.com>
Date: Sat, 20 Jul 2024 13:11:56 -0700
Subject: [PATCH] Fix cured villagers not having custom name removed

---
 src/main/java/dev/aurelium/auramobs/AuraMobs.java        | 9 +++++++--
 .../java/dev/aurelium/auramobs/listeners/MobSpawn.java   | 7 +------
 .../dev/aurelium/auramobs/listeners/MobTransform.java    | 2 +-
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/main/java/dev/aurelium/auramobs/AuraMobs.java b/src/main/java/dev/aurelium/auramobs/AuraMobs.java
index 53ae52a..6063e60 100644
--- a/src/main/java/dev/aurelium/auramobs/AuraMobs.java
+++ b/src/main/java/dev/aurelium/auramobs/AuraMobs.java
@@ -22,8 +22,7 @@
 import org.bstats.bukkit.Metrics;
 import org.bukkit.Bukkit;
 import org.bukkit.NamespacedKey;
-import org.bukkit.entity.LivingEntity;
-import org.bukkit.entity.Player;
+import org.bukkit.entity.*;
 import org.bukkit.persistence.PersistentDataType;
 import org.bukkit.plugin.java.JavaPlugin;
 
@@ -194,6 +193,12 @@ public boolean isAuraMob(LivingEntity m) {
         return m.getPersistentDataContainer().has(mobKey, PersistentDataType.INTEGER);
     }
 
+    public boolean isInvalidEntity(Entity entity) {
+        if (entity instanceof Boss || !(entity instanceof LivingEntity)) return true; // Types to exclude
+        if (entity instanceof Hoglin || entity instanceof Slime) return false; // Types to include
+        return !(entity instanceof Monster);
+    }
+
     public Formatter getFormatter() {
         return formatter;
     }
diff --git a/src/main/java/dev/aurelium/auramobs/listeners/MobSpawn.java b/src/main/java/dev/aurelium/auramobs/listeners/MobSpawn.java
index e285cbd..439c00c 100644
--- a/src/main/java/dev/aurelium/auramobs/listeners/MobSpawn.java
+++ b/src/main/java/dev/aurelium/auramobs/listeners/MobSpawn.java
@@ -32,7 +32,7 @@ public void onSpawn(CreatureSpawnEvent e) {
             }
             if (!f) return;
 
-            if (isInvalidEntity(e.getEntity())) {
+            if (plugin.isInvalidEntity(e.getEntity())) {
                 return;
             }
             LivingEntity entity = e.getEntity();
@@ -62,11 +62,6 @@ public void onSpawn(CreatureSpawnEvent e) {
         }
     }
 
-    private boolean isInvalidEntity(Entity entity) {
-        if (entity instanceof Boss || !(entity instanceof LivingEntity)) return true; // Types to exclude
-        if (entity instanceof Hoglin || entity instanceof Slime) return false; // Types to include
-        return !(entity instanceof Monster);
-    }
 
     public boolean passWorld(World world) {
         if (plugin.isWorldWhitelist()) {
diff --git a/src/main/java/dev/aurelium/auramobs/listeners/MobTransform.java b/src/main/java/dev/aurelium/auramobs/listeners/MobTransform.java
index db8583b..ae7af8a 100644
--- a/src/main/java/dev/aurelium/auramobs/listeners/MobTransform.java
+++ b/src/main/java/dev/aurelium/auramobs/listeners/MobTransform.java
@@ -27,7 +27,7 @@ public void onMutate(EntityTransformEvent event) {
             return;
         }
 
-        if (!(event.getTransformedEntity() instanceof LivingEntity)) {
+        if (plugin.isInvalidEntity(entity)) {
             event.getTransformedEntity().getPersistentDataContainer().remove(plugin.getMobKey());
             event.getTransformedEntity().setCustomNameVisible(false);
             event.getTransformedEntity().setCustomName(null);