Add before and after batch hooks

misode · Dec 25, 2023 · 3a6429f · 3a6429f
1 parent 4748957
commit 3a6429f
Show file tree

Hide file tree

Showing 4 changed files with 75 additions and 2 deletions.
diff --git a/src/main/java/io/github/misode/packtest/PackTestFunction.java b/src/main/java/io/github/misode/packtest/PackTestFunction.java
@@ -9,12 +9,12 @@
 import net.minecraft.commands.execution.ExecutionContext;
 import net.minecraft.commands.functions.CommandFunction;
 import net.minecraft.commands.functions.InstantiatedFunction;
-import net.minecraft.core.BlockPos;
 import net.minecraft.gametest.framework.GameTestHelper;
 import net.minecraft.gametest.framework.StructureUtils;
 import net.minecraft.gametest.framework.TestFunction;
 import net.minecraft.network.chat.CommonComponents;
 import net.minecraft.resources.ResourceLocation;
+import net.minecraft.server.level.ServerLevel;
 import net.minecraft.world.phys.Vec2;
 import net.minecraft.world.phys.Vec3;
 import org.jetbrains.annotations.Nullable;
@@ -107,6 +107,22 @@ private boolean isRequired() {
         return Optional.ofNullable(this.directives.get("optional")).map(s -> !Boolean.parseBoolean(s)).orElse(true);
     }
 
+    public void registerBatchHook(int permissionLevel, Map<String, Consumer<ServerLevel>> map, String type) {
+        String command = this.directives.get(type + "batch");
+        if (command == null) {
+            return;
+        }
+        String batchName = this.getBatchName();
+        Consumer<ServerLevel> oldBefore = map.putIfAbsent(batchName, (level) -> {
+            CommandSourceStack source = level.getServer().createCommandSourceStack()
+                    .withPermission(permissionLevel);
+            level.getServer().getCommands().performPrefixedCommand(source, command);
+        });
+        if (oldBefore != null) {
+            PackTest.LOGGER.error("Only one @" + type + "batch is allowed per batch. Batch '" + batchName + "' has more than one!");
+        }
+    }
+
     public TestFunction toTestFunction(int permissionLevel, CommandDispatcher<CommandSourceStack> dispatcher) {
         return new TestFunction(
                 this.getBatchName(),
@@ -148,7 +164,6 @@ private Consumer<GameTestHelper> createTestBody(int permissionLevel, CommandDisp
             Dummy dummy;
             if (dummyPos != null) {
                 try {
-                    PackTest.LOGGER.info("Directive position {} {}", dummyPos, helper.getLevel().getBlockState(BlockPos.containing(dummyPos.x, dummyPos.y, dummyPos.z)));
                     dummy = Dummy.createRandom(helper.getLevel().getServer(), helper.getLevel().dimension(), dummyPos);
                     dummy.setOnGround(true); // little hack because we know the dummy will be on the ground
                     sourceStack = sourceStack.withEntity(dummy);

diff --git a/src/main/java/io/github/misode/packtest/PackTestLibrary.java b/src/main/java/io/github/misode/packtest/PackTestLibrary.java
@@ -17,13 +17,15 @@
 import net.minecraft.util.profiling.ProfilerFiller;
 import net.minecraft.world.level.block.entity.StructureBlockEntity;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.io.BufferedReader;
 import java.io.IOException;
 import java.util.*;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.Executor;
+import java.util.function.Consumer;
 
 public class PackTestLibrary implements PreparableReloadListener {
     public static final PackTestLibrary INSTANCE = new PackTestLibrary(2, new CommandDispatcher<>());
@@ -33,6 +35,8 @@ public class PackTestLibrary implements PreparableReloadListener {
     private CommandDispatcher<CommandSourceStack> dispatcher;
     private Collection<TestFunction> tests = Lists.newArrayList();
     private Set<String> namespaces = Sets.newHashSet();
+    private Map<String, Consumer<ServerLevel>> beforeBatch = Maps.newHashMap();
+    private Map<String, Consumer<ServerLevel>> afterBatch = Maps.newHashMap();
     private final Map<String, GameTestHelper> testHelperMap = Maps.newHashMap();
 
     public PackTestLibrary(int permissionLevel, CommandDispatcher<CommandSourceStack> dispatcher) {
@@ -106,19 +110,25 @@ public Optional<GameTestHelper> getHelperAt(BlockPos pos, ServerLevel level) {
         });
 
         return prep.thenCompose(barrier::wait).thenAcceptAsync(functions -> {
+            Map<String, Consumer<ServerLevel>> beforeBatch = Maps.newHashMap();
+            Map<String, Consumer<ServerLevel>> afterBatch = Maps.newHashMap();
             ImmutableList.Builder<TestFunction> testsBuilder = ImmutableList.builder();
             ImmutableSet.Builder<String> namespacesBuilder = ImmutableSet.builder();
             functions.forEach((id, future) -> future.handle((val, err) -> {
                 if (err != null) {
                     PackTest.LOGGER.error("Failed to load test {}", id, err);
                 } else {
+                    val.registerBatchHook(this.permissionLevel, beforeBatch, "before");
+                    val.registerBatchHook(this.permissionLevel, afterBatch, "after");
                     testsBuilder.add(val.toTestFunction(this.permissionLevel, this.dispatcher));
                     namespacesBuilder.add(id.getNamespace());
                 }
                 return null;
             }).join());
             this.tests = testsBuilder.build();
             this.namespaces = namespacesBuilder.build();
+            this.beforeBatch = beforeBatch;
+            this.afterBatch = afterBatch;
             PackTest.LOGGER.info("Loaded {} tests", this.tests.size());
         });
     }
@@ -131,6 +141,14 @@ public Collection<String> getAllTestClassNames() {
         return namespaces;
     }
 
+    public @Nullable Consumer<ServerLevel> getBeforeBatchFunction(String batchName) {
+        return this.beforeBatch.get(batchName);
+    }
+
+    public @Nullable Consumer<ServerLevel> getAfterBatchFunction(String batchName) {
+        return this.afterBatch.get(batchName);
+    }
+
     private static List<String> readLines(Resource resource) {
         try (BufferedReader lvt1 = resource.openAsReader()) {
             return lvt1.lines().toList();

diff --git a/src/main/java/io/github/misode/packtest/mixin/GameTestBatchMixin.java b/src/main/java/io/github/misode/packtest/mixin/GameTestBatchMixin.java
@@ -0,0 +1,39 @@
+package io.github.misode.packtest.mixin;
+
+import io.github.misode.packtest.PackTestLibrary;
+import net.minecraft.gametest.framework.GameTestBatch;
+import net.minecraft.server.level.ServerLevel;
+import org.spongepowered.asm.mixin.Final;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Shadow;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+
+import java.util.function.Consumer;
+
+@Mixin(GameTestBatch.class)
+public class GameTestBatchMixin {
+
+    @Shadow
+    @Final
+    private String name;
+
+    @Inject(method = "runBeforeBatchFunction", at = @At("TAIL"))
+    private void runBeforeBatchFunction(ServerLevel level, CallbackInfo ci) {
+        String batchName = this.name.substring(0, this.name.lastIndexOf(":"));
+        Consumer<ServerLevel> consumer = PackTestLibrary.INSTANCE.getBeforeBatchFunction(batchName);
+        if (consumer != null) {
+            consumer.accept(level);
+        }
+    }
+
+    @Inject(method = "runAfterBatchFunction", at = @At("TAIL"))
+    private void runAfterBatchFunction(ServerLevel level, CallbackInfo ci) {
+        String batchName = this.name.substring(0, this.name.lastIndexOf(":"));
+        Consumer<ServerLevel> consumer = PackTestLibrary.INSTANCE.getAfterBatchFunction(batchName);
+        if (consumer != null) {
+            consumer.accept(level);
+        }
+    }
+}
diff --git a/src/main/resources/packtest.mixins.json b/src/main/resources/packtest.mixins.json
@@ -10,6 +10,7 @@
     "CommandsMixin",
     "EntityArgumentMixin",
     "EntitySelectorMixin",
+    "GameTestBatchMixin",
     "GameTestHelperMixin",
     "GameTestInfoMixin",
     "GameTestRegistryMixin",