Restructure transmitter chunk reloading tests so that the event listeners are only present for the test

Author: pupnewfster

src/gameTest/java/mekanism/common/tests/network/InventoryNetworkTest.java

@@ -1,26 +1,24 @@
 package mekanism.common.tests.network;
 
+import static mekanism.common.tests.util.TransporterTestUtils.colored;
+import static mekanism.common.tests.util.TransporterTestUtils.configured;
+import static mekanism.common.tests.util.TransporterTestUtils.containing;
+
 import java.util.function.Supplier;
-import mekanism.api.NBTConstants;
 import mekanism.api.text.EnumColor;
 import mekanism.api.tier.AlloyTier;
 import mekanism.common.lib.transmitter.ConnectionType;
 import mekanism.common.registries.MekanismBlocks;
 import mekanism.common.registries.MekanismItems;
 import mekanism.common.tests.MekanismTests;
 import mekanism.common.tests.util.GameTestUtils;
-import mekanism.common.util.EnumUtils;
-import mekanism.common.util.NBTUtils;
 import net.minecraft.SharedConstants;
 import net.minecraft.core.BlockPos;
 import net.minecraft.core.Direction;
 import net.minecraft.gametest.framework.GameTest;
-import net.minecraft.nbt.CompoundTag;
-import net.minecraft.nbt.ListTag;
 import net.minecraft.world.entity.player.Player;
 import net.minecraft.world.item.ItemStack;
 import net.minecraft.world.item.Items;
-import net.minecraft.world.level.ItemLike;
 import net.minecraft.world.level.block.Blocks;
 import net.minecraft.world.level.levelgen.structure.templatesystem.StructureTemplate;
 import net.neoforged.testframework.DynamicTest;
@@ -29,7 +27,6 @@
 import net.neoforged.testframework.annotation.TestHolder;
 import net.neoforged.testframework.gametest.ExtendedGameTestHelper;
 import net.neoforged.testframework.gametest.StructureTemplateBuilder;
-import org.jetbrains.annotations.Nullable;
 
 @ForEachTest(groups = "network.inventory")
 public class InventoryNetworkTest {
@@ -56,64 +53,6 @@ public class InventoryNetworkTest {
           .fill(9, 0, 0, 9, 0, 1, MekanismBlocks.BASIC_LOGISTICAL_TRANSPORTER.getBlock().defaultBlockState())
     );
 
-
-    private static CompoundTag containing(ItemLike itemLike) {
-        return containing(new ItemStack(itemLike));
-    }
-
-    private static CompoundTag containing(ItemLike itemLike, int amount) {
-        return containing(new ItemStack(itemLike, amount));
-    }
-
-    private static CompoundTag containing(ItemStack... stacks) {
-        CompoundTag tag = new CompoundTag();
-        ListTag items = new ListTag();
-        for (int i = 0; i < stacks.length; i++) {
-            CompoundTag item = stacks[i].save(new CompoundTag());
-            item.putByte(NBTConstants.SLOT, (byte) i);
-            items.add(item);
-        }
-        tag.put(NBTConstants.ITEMS, items);
-        return tag;
-    }
-
-    @Nullable
-    private static CompoundTag colored(EnumColor color) {
-        return colored(color, null);
-    }
-
-    @Nullable
-    private static CompoundTag colored(@Nullable EnumColor color, @Nullable Direction pull) {
-        return configured(color, pull, ConnectionType.PULL);
-    }
-
-    @Nullable
-    private static CompoundTag configured(Direction side) {
-        return configured(side, ConnectionType.PULL);
-    }
-
-    @Nullable
-    private static CompoundTag configured(Direction side, ConnectionType connectionType) {
-        return configured(null, side, connectionType);
-    }
-
-    @Nullable
-    private static CompoundTag configured(@Nullable EnumColor color, @Nullable Direction side, ConnectionType connectionType) {
-        if (color == null && side == null) {
-            return null;
-        }
-        CompoundTag tag = new CompoundTag();
-        if (color != null) {
-            NBTUtils.writeEnum(tag, NBTConstants.COLOR, color);
-        }
-        if (side != null) {
-            int[] raw = new int[EnumUtils.DIRECTIONS.length];
-            raw[side.ordinal()] = connectionType.ordinal();
-            tag.putIntArray(NBTConstants.CONNECTION, raw);
-        }
-        return tag;
-    }
-
     //TODO: Do we want to somehow test the case of when we add a shorter path to the same destination, as the newly pulled items go via the new shorter path
     // but any already en-route ones continue the original way
 

src/gameTest/java/mekanism/common/tests/network/TransmitterNetworkTest.java

@@ -2,25 +2,33 @@
 
 import java.util.UUID;
 import java.util.function.Supplier;
-import mekanism.api.functions.TriConsumer;
+import mekanism.api.functions.ConstantPredicates;
+import mekanism.common.Mekanism;
 import mekanism.common.content.network.transmitter.Transmitter;
 import mekanism.common.lib.transmitter.DynamicNetwork;
 import mekanism.common.registries.MekanismBlocks;
 import mekanism.common.tests.MekanismTests;
 import mekanism.common.tests.util.GameTestUtils;
 import mekanism.common.tile.transmitter.TileEntityTransmitter;
+import mekanism.common.util.WorldUtils;
 import net.minecraft.SharedConstants;
 import net.minecraft.core.BlockPos;
 import net.minecraft.gametest.framework.GameTest;
+import net.minecraft.server.level.ChunkHolder;
+import net.minecraft.server.level.ChunkLevel;
+import net.minecraft.server.level.ChunkMap;
+import net.minecraft.server.level.DistanceManager;
+import net.minecraft.server.level.ServerLevel;
 import net.minecraft.world.level.ChunkPos;
 import net.minecraft.world.level.block.Blocks;
 import net.minecraft.world.level.levelgen.structure.templatesystem.StructureTemplate;
+import net.neoforged.neoforge.event.level.ChunkEvent;
+import net.neoforged.testframework.DynamicTest;
 import net.neoforged.testframework.annotation.ForEachTest;
 import net.neoforged.testframework.annotation.RegisterStructureTemplate;
 import net.neoforged.testframework.annotation.TestHolder;
 import net.neoforged.testframework.gametest.ExtendedGameTestHelper;
 import net.neoforged.testframework.gametest.StructureTemplateBuilder;
-import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 @ForEachTest(groups = "network.transmitter")
@@ -38,8 +46,22 @@ public class TransmitterNetworkTest {
 
     @GameTest(template = STRAIGHT_CABLE, setupTicks = SETUP_TICKS, batch = "1")
     @TestHolder(description = "Tests that reloading intermediary chunks does not cause a network to break.")
-    public static void reloadIntermediary(final ExtendedGameTestHelper helper) {
-        GameTestUtils.succeedIfAfterReload(helper, new ChunkPos(1, 0), new MatchingNetworkValidator(helper));
+    public static void reloadIntermediary(final DynamicTest test) {
+        final ChunkData chunkData = new ChunkData(1, 0);
+        test.eventListeners().forge().addListener((final ChunkEvent.Unload event) -> chunkData.updateChunk(event, false));
+        test.eventListeners().forge().addListener((final ChunkEvent.Load event) -> chunkData.updateChunk(event, true));
+
+        test.onGameTest(helper -> helper.startSequence()
+              .thenMap(() -> chunkData.updateChunkLoading(helper, false, GameTestUtils.UNLOAD_LEVEL))
+              .thenWaitUntil(() -> chunkData.waitFor(helper, false))
+              //Wait 5 ticks in case anything needs more time to process after the chunk unloads
+              .thenExecuteAfter(5, level -> chunkData.updateChunkLoading(helper, true, level))
+              .thenWaitUntil(() -> chunkData.waitFor(helper, true))
+              //Wait 5 ticks in case anything needs more time to process after the chunk loads
+              .thenIdle(5)
+              .thenExecute(new MatchingNetworkValidator(helper))
+              .thenSucceed()
+        );
     }
 
     /**
@@ -68,40 +90,10 @@ public static void inaccessibleNotUnloaded(final ExtendedGameTestHelper helper)
               .thenExecuteAfter(5, level -> GameTestUtils.setChunkLoadLevel(helper, relativeChunk, level))
               //Wait 5 ticks in case anything needs more time to process after the chunk loads
               .thenIdle(5)
-              .thenWaitUntil(0, new MatchingNetworkValidator(helper))
+              .thenExecute(new MatchingNetworkValidator(helper))
               .thenSucceed();
     }
 
-    private static void forEachTransmitter(ExtendedGameTestHelper helper, TriConsumer<TileEntityTransmitter, Transmitter<?, ?, ?>, BlockPos> consumer) {
-        forEachTransmitter(helper, true, consumer);
-    }
-
-    private static void forEachTransmitter(ExtendedGameTestHelper helper, boolean expectNetwork, TriConsumer<TileEntityTransmitter, Transmitter<?, ?, ?>, BlockPos> consumer) {
-        helper.forEveryBlockInStructure(relativePos -> {
-            TileEntityTransmitter blockEntity = getTransmitterNNAt(helper, relativePos);
-            Transmitter<?, ?, ?> transmitter = blockEntity.getTransmitter();
-            if (expectNetwork && !transmitter.hasTransmitterNetwork()) {
-                helper.fail("No transmitter network found", relativePos);
-            }
-            consumer.accept(blockEntity, transmitter, relativePos);
-        });
-    }
-
-    @Nullable
-    private static TileEntityTransmitter getTransmitterAt(ExtendedGameTestHelper helper, BlockPos relativePos) {
-        return GameTestUtils.getBlockEntity(helper, TileEntityTransmitter.class, relativePos);
-    }
-
-    @NotNull
-    private static TileEntityTransmitter getTransmitterNNAt(ExtendedGameTestHelper helper, BlockPos relativePos) {
-        TileEntityTransmitter transmitter = getTransmitterAt(helper, relativePos);
-        if (transmitter == null) {
-            helper.fail("Expected transmitter", relativePos);
-        }
-        //noinspection ConstantConditions (can't get heve if null as helper#fail throws an exception)
-        return transmitter;
-    }
-
     private static class MatchingNetworkValidator implements Runnable {
 
         private final ExtendedGameTestHelper helper;
@@ -113,14 +105,106 @@ public MatchingNetworkValidator(ExtendedGameTestHelper helper) {
 
         @Override
         public void run() {
-            forEachTransmitter(helper, (tile, transmitter, relativePos) -> {
-                DynamicNetwork<?, ?, ?> network = transmitter.getTransmitterNetwork();
-                if (networkUUID == null) {
-                    networkUUID = network.getUUID();
-                } else if (!networkUUID.equals(network.getUUID())) {
-                    helper.fail("Multiple transmitter networks", relativePos);
+            helper.forEveryBlockInStructure(relativePos -> {
+                if (WorldUtils.isBlockLoaded(helper.getLevel(), helper.absolutePos(relativePos))) {
+                    Transmitter<?, ?, ?> transmitter = helper.requireBlockEntity(relativePos, TileEntityTransmitter.class).getTransmitter();
+                    if (!transmitter.hasTransmitterNetwork()) {
+                        helper.fail("No transmitter network found", relativePos);
+                    }
+                    DynamicNetwork<?, ?, ?> network = transmitter.getTransmitterNetwork();
+                    if (networkUUID == null) {
+                        networkUUID = network.getUUID();
+                    } else if (!networkUUID.equals(network.getUUID())) {
+                        helper.fail("Multiple transmitter networks", relativePos);
+                    }
+                } else {
+                    helper.fail("Expected expected position to be loaded", relativePos);
                 }
             });
         }
     }
+
+    private static class ChunkData {
+
+        private static final boolean DEBUG_CHUNK_LOADING = false;
+
+        private final ChunkPos relativePos;
+        @Nullable
+        private ChunkPos absolutePos;
+        private long absPos;
+        private boolean isLoaded;
+
+        public ChunkData(int x, int y) {
+            this.relativePos = new ChunkPos(x, y);
+        }
+
+        //TODO - GameTest: Can we make unloads not cause the game to crash if a player tries to run them in world? It crashes as we force unload a chunk that would be near the player and it crashes from a null pointer
+        public int updateChunkLoading(ExtendedGameTestHelper helper, boolean load, int level) {
+            if (absolutePos == null) {
+                absolutePos = GameTestUtils.absolutePos(helper, relativePos);
+                absPos = absolutePos.toLong();
+            }
+            ServerLevel serverLevel = helper.getLevel();
+            if (WorldUtils.isChunkLoaded(serverLevel, absolutePos) != load) {
+                //If the chunk isn't watched and is loaded we want to try and unload it
+                ChunkMap chunkMap = serverLevel.getChunkSource().chunkMap;
+                DistanceManager distanceManager = chunkMap.getDistanceManager();
+                ChunkHolder holder = distanceManager.getChunk(absPos);
+                //Watch the chunk and mark whether it is currently loaded or not
+                if ((holder == null) == load) {
+                    //If it is loaded then we need to try and unload it
+                    isLoaded = !load;
+                    if (DEBUG_CHUNK_LOADING) {
+                        Mekanism.logger.info("Trying to {} chunk at: {}", load ? "load" : "unload", absolutePos);
+                    }
+                    if (load) {
+                        //Load the chunk to the level it was unloaded at
+                        holder = distanceManager.updateChunkScheduling(absPos, level, holder, GameTestUtils.UNLOAD_LEVEL);
+                        if (holder == null) {//Should never happen unless start value was unloaded
+                            fail(helper, "Error loading chunk");
+                        } else {
+                            //And ensure we schedule it based on the status (in general this should be ChunkStatus.FULL)
+                            chunkMap.schedule(holder, ChunkLevel.generationStatus(holder.getTicketLevel()));
+                        }
+                    } else {
+                        //If it is currently loaded, queue it for unload
+                        level = holder.getTicketLevel();
+                        distanceManager.updateChunkScheduling(absPos, GameTestUtils.UNLOAD_LEVEL, holder, level);
+                        //And then unload it
+                        chunkMap.processUnloads(ConstantPredicates.ALWAYS_TRUE);
+                    }
+                } else if (DEBUG_CHUNK_LOADING) {
+                    //Note: Even with debug logging enabled odds are this case isn't even possible due to the earlier check to skip if unloaded
+                    Mekanism.logger.info("Trying to {} already {} chunk at: {}", load ? "load" : "unload", load ? "loaded" : "unloaded", absolutePos);
+                }
+            } else if (DEBUG_CHUNK_LOADING) {
+                Mekanism.logger.info("Chunk at: {} is already {}", absolutePos, load ? "unloaded" : "loaded");
+            }
+            return level;
+        }
+
+        public void updateChunk(ChunkEvent event, boolean loaded) {
+            if (!event.getLevel().isClientSide() && event.getChunk().getPos().equals(absolutePos)) {
+                //If we are watching the chunk and the loaded state isn't what we already had it as
+                if (isLoaded != loaded) {
+                    isLoaded = loaded;
+                    if (DEBUG_CHUNK_LOADING) {
+                        Mekanism.logger.info("Chunk {}: {}", loaded ? "loaded" : "unloaded", absolutePos);
+                    }
+                } else if (DEBUG_CHUNK_LOADING) {
+                    Mekanism.logger.info("Chunk was already {}: {}", loaded ? "loaded" : "unloaded", absolutePos);
+                }
+            }
+        }
+
+        public void fail(ExtendedGameTestHelper helper, String message) {
+            helper.fail(message + " at " + absolutePos + " (relative: " + relativePos + ")");
+        }
+
+        public void waitFor(ExtendedGameTestHelper helper, boolean loaded) {
+            if (isLoaded != loaded) {//If our loaded status does not match our desired one, keep throwing an exception until it does
+                fail(helper, "Chunk has not been marked as " + (loaded ? "loaded" : "unloaded") + " yet");
+            }
+        }
+    }
 }