Add some more tests (#413)

* More work on tests

* Update LivingEntityEventTests.java

* Test

* Updates

testframework/src/main/java/net/neoforged/testframework/gametest/ExtendedGameTestHelper.java

@@ -5,24 +5,30 @@
 
 package net.neoforged.testframework.gametest;
 
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.mojang.authlib.GameProfile;
 import io.netty.channel.embedded.EmbeddedChannel;
 import java.util.List;
+import java.util.Set;
 import java.util.UUID;
 import java.util.function.BiPredicate;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
 import net.minecraft.core.BlockPos;
 import net.minecraft.core.Direction;
+import net.minecraft.core.registries.BuiltInRegistries;
 import net.minecraft.gametest.framework.GameTestAssertException;
 import net.minecraft.gametest.framework.GameTestAssertPosException;
 import net.minecraft.gametest.framework.GameTestHelper;
 import net.minecraft.gametest.framework.GameTestInfo;
+import net.minecraft.gametest.framework.GameTestListener;
 import net.minecraft.network.Connection;
 import net.minecraft.network.ConnectionProtocol;
 import net.minecraft.network.protocol.PacketFlow;
 import net.minecraft.server.network.CommonListenerCookie;
+import net.minecraft.world.Difficulty;
 import net.minecraft.world.InteractionHand;
 import net.minecraft.world.InteractionResult;
 import net.minecraft.world.effect.MobEffect;
@@ -42,6 +48,9 @@
 import net.minecraft.world.phys.AABB;
 import net.minecraft.world.phys.BlockHitResult;
 import net.minecraft.world.phys.Vec3;
+import net.neoforged.bus.api.Event;
+import net.neoforged.neoforge.common.NeoForge;
+import net.neoforged.neoforge.event.entity.living.LivingKnockBackEvent;
 import org.jetbrains.annotations.Nullable;
 
 public class ExtendedGameTestHelper extends GameTestHelper {
@@ -109,6 +118,7 @@ public void tick() {
         this.getLevel().getServer().getConnection().getConnections().add(connection);
         this.testInfo.addListener(serverplayer);
         serverplayer.gameMode.changeGameModeForPlayer(gameType);
+        serverplayer.setYRot(180);
         return serverplayer;
     }
 
@@ -203,11 +213,93 @@ public void pulseRedstone(int x, int y, int z, long delay) {
         pulseRedstone(new BlockPos(x, y, z), delay);
     }
 
+    public void assertPlayerHasItem(Player player, Item item) {
+        assertTrue(player.getInventory().hasAnyOf(Set.of(item)), "Player doesn't have '" + BuiltInRegistries.ITEM.getKey(item) + "' in their inventory!");
+    }
+
+    public void requireDifficulty(final Difficulty difficulty) {
+        final var oldDifficulty = getLevel().getServer().getWorldData().getDifficulty();
+        if (oldDifficulty != difficulty) {
+            getLevel().getServer().setDifficulty(difficulty, true);
+            addEndListener(passed -> getLevel().getServer().setDifficulty(oldDifficulty, true));
+        }
+    }
+
+    public void addEndListener(Consumer<Boolean> listener) {
+        testInfo.addListener(new GameTestListener() {
+            @Override
+            public void testStructureLoaded(GameTestInfo info) {
+
+            }
+
+            @Override
+            public void testPassed(GameTestInfo info) {
+                listener.accept(true);
+            }
+
+            @Override
+            public void testFailed(GameTestInfo info) {
+                listener.accept(false);
+            }
+        });
+    }
+
+    public <T> T catchException(final ThrowingSupplier<T> supplier) {
+        try {
+            return supplier.get();
+        } catch (Throwable throwable) {
+            throw new GameTestAssertException(throwable.getMessage());
+        }
+    }
+
+    public void catchException(final ThrowingRunnable run) {
+        try {
+            run.run();
+        } catch (Throwable throwable) {
+            throw new GameTestAssertException(throwable.getMessage());
+        }
+    }
+
+    public <T extends Entity> T requireEntityAt(EntityType<T> type, int x, int y, int z) {
+        return requireEntityAt(type, new BlockPos(x, y, z));
+    }
+
+    public <T extends Entity> T requireEntityAt(EntityType<T> type, BlockPos pos) {
+        final var inRange = getEntities(type, pos, 2);
+        assertTrue(inRange.size() == 1, "Only one entity must be present at " + pos);
+        return inRange.get(0);
+    }
+
+    @CanIgnoreReturnValue
+    public <T extends Entity> T knockbackResistant(T entity) {
+        addTemporaryListener((final LivingKnockBackEvent event) -> {
+            if (event.getEntity().getUUID().equals(entity.getUUID())) {
+                event.setCanceled(true);
+            }
+        });
+        return entity;
+    }
+
+    public <T extends Event> void addTemporaryListener(Consumer<T> event) {
+        NeoForge.EVENT_BUS.addListener(event);
+        addEndListener(success -> NeoForge.EVENT_BUS.unregister(event));
+    }
+
     public <E extends LivingEntity> void assertMobEffectPresent(E entity, MobEffect effect, String testName) {
         assertEntityProperty(entity, e -> e.hasEffect(effect), testName);
     }
 
     public <E extends LivingEntity> void assertMobEffectAbsent(E entity, MobEffect effect, String testName) {
         assertEntityProperty(entity, e -> !e.hasEffect(effect), testName);
     }
+
+    @FunctionalInterface
+    public interface ThrowingSupplier<T> {
+        T get() throws Throwable;
+    }
+
+    @FunctionalInterface
+    public interface ThrowingRunnable {
+        void run() throws Throwable;
+    }
 }

testframework/src/main/java/net/neoforged/testframework/gametest/ParametrizedGameTestSequence.java

@@ -11,6 +11,7 @@
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
+import net.minecraft.gametest.framework.GameTestAssertException;
 import net.minecraft.gametest.framework.GameTestInfo;
 import net.minecraft.gametest.framework.GameTestSequence;
 
@@ -25,7 +26,13 @@ public ParametrizedGameTestSequence(GameTestInfo info, ExtendedSequence sequence
 
         final AtomicReference<T> val = new AtomicReference<>();
         sequence.thenExecute(() -> val.set(value.get()));
-        this.value = val::get;
+        this.value = () -> {
+            final var v = val.get();
+            if (v == null) {
+                throw new GameTestAssertException("Expected value to be non-null!");
+            }
+            return v;
+        };
     }
 
     public ParametrizedGameTestSequence<T> thenWaitUntil(Runnable condition) {

testframework/src/main/java/net/neoforged/testframework/impl/reg/RegistrationHelperImpl.java

@@ -13,8 +13,10 @@
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import net.minecraft.core.Registry;
 import net.minecraft.core.registries.Registries;
 import net.minecraft.data.CachedOutput;
@@ -25,6 +27,7 @@
 import net.minecraft.resources.ResourceLocation;
 import net.neoforged.bus.api.Event;
 import net.neoforged.bus.api.IEventBus;
+import net.neoforged.neoforge.attachment.AttachmentType;
 import net.neoforged.neoforge.client.model.generators.BlockStateProvider;
 import net.neoforged.neoforge.client.model.generators.ItemModelProvider;
 import net.neoforged.neoforge.client.model.generators.ModelProvider;
@@ -33,6 +36,7 @@
 import net.neoforged.neoforge.common.data.LanguageProvider;
 import net.neoforged.neoforge.data.event.GatherDataEvent;
 import net.neoforged.neoforge.registries.DeferredRegister;
+import net.neoforged.testframework.registration.DeferredAttachmentTypes;
 import net.neoforged.testframework.registration.DeferredBlocks;
 import net.neoforged.testframework.registration.DeferredEntityTypes;
 import net.neoforged.testframework.registration.DeferredItems;
@@ -137,6 +141,13 @@ public DeferredEntityTypes entityTypes() {
         return entityTypes;
     }
 
+    private final DeferredRegistrar<AttachmentType<?>, DeferredAttachmentTypes> attachments = new DeferredRegistrar<>((namespace, reg) -> new DeferredAttachmentTypes(namespace));
+
+    @Override
+    public DeferredAttachmentTypes attachments() {
+        return attachments.get();
+    }
+
     @Override
     public String modId() {
         return modId;
@@ -187,4 +198,24 @@ public String getName() {
             }
         }));
     }
+
+    private final class DeferredRegistrar<R, T extends DeferredRegister<R>> implements Supplier<T> {
+        private final BiFunction<String, RegistrationHelper, T> factory;
+
+        private T cached;
+
+        private DeferredRegistrar(BiFunction<String, RegistrationHelper, T> factory) {
+            this.factory = factory;
+        }
+
+        @Override
+        public T get() {
+            if (cached == null) {
+                cached = factory.apply(RegistrationHelperImpl.this.modId, RegistrationHelperImpl.this);
+                registrars.put(cached.getRegistryKey(), cached);
+                if (RegistrationHelperImpl.this.bus != null) cached.register(bus);
+            }
+            return cached;
+        }
+    }
 }

testframework/src/main/java/net/neoforged/testframework/registration/DeferredAttachmentTypes.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright (c) NeoForged and contributors
+ * SPDX-License-Identifier: LGPL-2.1-only
+ */
+
+package net.neoforged.testframework.registration;
+
+import java.util.function.Supplier;
+import java.util.function.UnaryOperator;
+import net.neoforged.neoforge.attachment.AttachmentType;
+import net.neoforged.neoforge.registries.DeferredRegister;
+import net.neoforged.neoforge.registries.NeoForgeRegistries;
+
+public class DeferredAttachmentTypes extends DeferredRegister<AttachmentType<?>> {
+    public DeferredAttachmentTypes(String namespace) {
+        super(NeoForgeRegistries.Keys.ATTACHMENT_TYPES, namespace);
+    }
+
+    public <T> AttachmentType<T> registerSimpleAttachment(String name, Supplier<T> defaultValue) {
+        return register(name, defaultValue, UnaryOperator.identity());
+    }
+
+    public <T> AttachmentType<T> register(String name, Supplier<T> defaultValue, UnaryOperator<AttachmentType.Builder<T>> factory) {
+        final var attach = factory.apply(AttachmentType.builder(defaultValue)).build();
+        register(name, () -> attach);
+        return attach;
+    }
+}

testframework/src/main/java/net/neoforged/testframework/registration/RegistrationHelper.java

@@ -40,6 +40,11 @@ public interface RegistrationHelper {
      */
     DeferredEntityTypes entityTypes();
 
+    /**
+     * {@return a helper for attachment type registration}
+     */
+    DeferredAttachmentTypes attachments();
+
     /**
      * {@return the mod id of this helper}
      */

tests/src/main/java/net/neoforged/neoforge/debug/client/ClientEventTests.java

@@ -9,6 +9,7 @@
 import net.minecraft.network.chat.Component;
 import net.neoforged.api.distmarker.Dist;
 import net.neoforged.neoforge.client.event.ClientChatEvent;
+import net.neoforged.neoforge.client.event.ClientPlayerChangeGameTypeEvent;
 import net.neoforged.testframework.DynamicTest;
 import net.neoforged.testframework.annotation.ForEachTest;
 import net.neoforged.testframework.annotation.TestHolder;
@@ -25,4 +26,9 @@ static void playerClientChatEvent(final ClientChatEvent event, final DynamicTest
             Minecraft.getInstance().tell(() -> test.requestConfirmation(Minecraft.getInstance().player, Component.literal("Was your message modified?")));
         }
     }
+
+    @TestHolder(description = { "Tests if the ClientPlayerChangeGameTypeEvent event is fired", "Will ask the player for confirmation when the player changes their gamemode" })
+    static void clientPlayerChangeGameTypeEvent(final ClientPlayerChangeGameTypeEvent event, final DynamicTest test) {
+        test.requestConfirmation(Minecraft.getInstance().player, Component.literal("Did you just change your game mode from " + event.getCurrentGameType() + " to " + event.getNewGameType() + "?"));
+    }
 }

tests/src/main/java/net/neoforged/neoforge/debug/entity/EntityEventTests.java

@@ -5,17 +5,21 @@
 
 package net.neoforged.neoforge.debug.entity;
 
+import net.minecraft.core.registries.Registries;
 import net.minecraft.gametest.framework.GameTest;
 import net.minecraft.world.entity.EntityType;
+import net.minecraft.world.entity.ai.attributes.RangedAttribute;
 import net.minecraft.world.entity.animal.Cow;
 import net.minecraft.world.entity.animal.Pig;
 import net.minecraft.world.item.ItemStack;
 import net.minecraft.world.item.Items;
+import net.neoforged.neoforge.event.entity.EntityAttributeModificationEvent;
 import net.neoforged.neoforge.event.entity.EntityTeleportEvent;
 import net.neoforged.testframework.DynamicTest;
 import net.neoforged.testframework.annotation.ForEachTest;
 import net.neoforged.testframework.annotation.TestHolder;
 import net.neoforged.testframework.gametest.EmptyTemplate;
+import net.neoforged.testframework.registration.RegistrationHelper;
 
 @ForEachTest(groups = { EntityTests.GROUP + ".event", "event" })
 public class EntityEventTests {
@@ -50,4 +54,19 @@ static void entityTeleportEvent(final DynamicTest test) {
                 .thenExecute(helper::killAllEntities)
                 .thenSucceed());
     }
+
+    @GameTest
+    @EmptyTemplate(floor = true)
+    @TestHolder(description = "Tests if the EntityAttributeModificationEvent is fired")
+    static void entityAttributeModificationEvent(final DynamicTest test, final RegistrationHelper reg) {
+        final var testAttr = reg.registrar(Registries.ATTRIBUTE).register("test_attribute", () -> new RangedAttribute(reg.modId() + ".test_attr", 1.5D, 0.0D, 1024.0D).setSyncable(true));
+        test.framework().modEventBus().addListener((final EntityAttributeModificationEvent event) -> {
+            event.add(EntityType.DONKEY, testAttr.get());
+        });
+
+        test.onGameTest(helper -> helper.startSequence(() -> helper.spawnWithNoFreeWill(EntityType.DONKEY, 1, 2, 1))
+                .thenExecute(donkey -> helper.assertEntityProperty(
+                        donkey, d -> d.getAttribute(testAttr.get()).getValue(), "test attribute", 1.5D))
+                .thenSucceed());
+    }
 }