minecraft-src/net/minecraft/gametest/framework/TestEnvironmentDefinition.java

package net.minecraft.gametest.framework;

import com.mojang.logging.LogUtils;
import com.mojang.serialization.Codec;
import com.mojang.serialization.MapCodec;
import com.mojang.serialization.codecs.RecordCodecBuilder;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import net.minecraft.commands.CommandSourceStack;
import net.minecraft.commands.functions.CommandFunction;
import net.minecraft.core.Holder;
import net.minecraft.core.Registry;
import net.minecraft.core.registries.BuiltInRegistries;
import net.minecraft.core.registries.Registries;
import net.minecraft.gametest.framework.TestEnvironmentDefinition.SetGameRules.Entry;
import net.minecraft.gametest.framework.TestEnvironmentDefinition.Weather.Type;
import net.minecraft.resources.RegistryFileCodec;
import net.minecraft.resources.ResourceLocation;
import net.minecraft.server.MinecraftServer;
import net.minecraft.server.ServerFunctionManager;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.util.ExtraCodecs;
import net.minecraft.world.level.GameRules;
import org.slf4j.Logger;

public interface TestEnvironmentDefinition {
	Codec<TestEnvironmentDefinition> DIRECT_CODEC = BuiltInRegistries.TEST_ENVIRONMENT_DEFINITION_TYPE
		.byNameCodec()
		.dispatch(TestEnvironmentDefinition::codec, mapCodec -> mapCodec);
	Codec<Holder<TestEnvironmentDefinition>> CODEC = RegistryFileCodec.create(Registries.TEST_ENVIRONMENT, DIRECT_CODEC);

	static MapCodec<? extends TestEnvironmentDefinition> bootstrap(Registry<MapCodec<? extends TestEnvironmentDefinition>> registry) {
		Registry.register(registry, "all_of", TestEnvironmentDefinition.AllOf.CODEC);
		Registry.register(registry, "game_rules", TestEnvironmentDefinition.SetGameRules.CODEC);
		Registry.register(registry, "time_of_day", TestEnvironmentDefinition.TimeOfDay.CODEC);
		Registry.register(registry, "weather", TestEnvironmentDefinition.Weather.CODEC);
		return Registry.register(registry, "function", TestEnvironmentDefinition.Functions.CODEC);
	}

	void setup(ServerLevel level);

	default void teardown(ServerLevel level) {
	}

	MapCodec<? extends TestEnvironmentDefinition> codec();

	public record AllOf(List<Holder<TestEnvironmentDefinition>> definitions) implements TestEnvironmentDefinition {
		public static final MapCodec<TestEnvironmentDefinition.AllOf> CODEC = RecordCodecBuilder.mapCodec(
			instance -> instance.group(TestEnvironmentDefinition.CODEC.listOf().fieldOf("definitions").forGetter(TestEnvironmentDefinition.AllOf::definitions))
				.apply(instance, TestEnvironmentDefinition.AllOf::new)
		);

		public AllOf(TestEnvironmentDefinition... definitons) {
			this(Arrays.stream(definitons).map(Holder::direct).toList());
		}

		@Override
		public void setup(ServerLevel level) {
			this.definitions.forEach(holder -> ((TestEnvironmentDefinition)holder.value()).setup(level));
		}

		@Override
		public void teardown(ServerLevel level) {
			this.definitions.forEach(holder -> ((TestEnvironmentDefinition)holder.value()).teardown(level));
		}

		@Override
		public MapCodec<TestEnvironmentDefinition.AllOf> codec() {
			return CODEC;
		}
	}

	public record Functions(Optional<ResourceLocation> setupFunction, Optional<ResourceLocation> teardownFunction) implements TestEnvironmentDefinition {
		private static final Logger LOGGER = LogUtils.getLogger();
		public static final MapCodec<TestEnvironmentDefinition.Functions> CODEC = RecordCodecBuilder.mapCodec(
			instance -> instance.group(
					ResourceLocation.CODEC.optionalFieldOf("setup").forGetter(TestEnvironmentDefinition.Functions::setupFunction),
					ResourceLocation.CODEC.optionalFieldOf("teardown").forGetter(TestEnvironmentDefinition.Functions::teardownFunction)
				)
				.apply(instance, TestEnvironmentDefinition.Functions::new)
		);

		@Override
		public void setup(ServerLevel level) {
			this.setupFunction.ifPresent(resourceLocation -> run(level, resourceLocation));
		}

		@Override
		public void teardown(ServerLevel level) {
			this.teardownFunction.ifPresent(resourceLocation -> run(level, resourceLocation));
		}

		private static void run(ServerLevel level, ResourceLocation function) {
			MinecraftServer minecraftServer = level.getServer();
			ServerFunctionManager serverFunctionManager = minecraftServer.getFunctions();
			Optional<CommandFunction<CommandSourceStack>> optional = serverFunctionManager.get(function);
			if (optional.isPresent()) {
				CommandSourceStack commandSourceStack = minecraftServer.createCommandSourceStack().withPermission(2).withSuppressedOutput().withLevel(level);
				serverFunctionManager.execute((CommandFunction<CommandSourceStack>)optional.get(), commandSourceStack);
			} else {
				LOGGER.error("Test Batch failed for non-existent function {}", function);
			}
		}

		@Override
		public MapCodec<TestEnvironmentDefinition.Functions> codec() {
			return CODEC;
		}
	}

	public record SetGameRules(List<Entry<Boolean, GameRules.BooleanValue>> boolRules, List<Entry<Integer, GameRules.IntegerValue>> intRules)
		implements TestEnvironmentDefinition {
		public static final MapCodec<TestEnvironmentDefinition.SetGameRules> CODEC = RecordCodecBuilder.mapCodec(
			instance -> instance.group(
					Entry.codec(GameRules.BooleanValue.class, Codec.BOOL).listOf().fieldOf("bool_rules").forGetter(TestEnvironmentDefinition.SetGameRules::boolRules),
					Entry.codec(GameRules.IntegerValue.class, Codec.INT).listOf().fieldOf("int_rules").forGetter(TestEnvironmentDefinition.SetGameRules::intRules)
				)
				.apply(instance, TestEnvironmentDefinition.SetGameRules::new)
		);

		@Override
		public void setup(ServerLevel level) {
			GameRules gameRules = level.getGameRules();
			MinecraftServer minecraftServer = level.getServer();

			for (Entry<Boolean, GameRules.BooleanValue> entry : this.boolRules) {
				gameRules.<GameRules.BooleanValue>getRule(entry.key()).set((Boolean)entry.value(), minecraftServer);
			}

			for (Entry<Integer, GameRules.IntegerValue> entry : this.intRules) {
				gameRules.<GameRules.IntegerValue>getRule(entry.key()).set((Integer)entry.value(), minecraftServer);
			}
		}

		@Override
		public void teardown(ServerLevel level) {
			GameRules gameRules = level.getGameRules();
			MinecraftServer minecraftServer = level.getServer();

			for (Entry<Boolean, GameRules.BooleanValue> entry : this.boolRules) {
				gameRules.<GameRules.BooleanValue>getRule(entry.key()).setFrom(GameRules.<GameRules.BooleanValue>getType(entry.key()).createRule(), minecraftServer);
			}

			for (Entry<Integer, GameRules.IntegerValue> entry : this.intRules) {
				gameRules.<GameRules.IntegerValue>getRule(entry.key()).setFrom(GameRules.<GameRules.IntegerValue>getType(entry.key()).createRule(), minecraftServer);
			}
		}

		@Override
		public MapCodec<TestEnvironmentDefinition.SetGameRules> codec() {
			return CODEC;
		}

		public static <S, T extends GameRules.Value<T>> Entry<S, T> entry(GameRules.Key<T> key, S value) {
			return new Entry(key, value);
		}
	}

	public record TimeOfDay(int time) implements TestEnvironmentDefinition {
		public static final MapCodec<TestEnvironmentDefinition.TimeOfDay> CODEC = RecordCodecBuilder.mapCodec(
			instance -> instance.group(ExtraCodecs.NON_NEGATIVE_INT.fieldOf("time").forGetter(TestEnvironmentDefinition.TimeOfDay::time))
				.apply(instance, TestEnvironmentDefinition.TimeOfDay::new)
		);

		@Override
		public void setup(ServerLevel level) {
			level.setDayTime(this.time);
		}

		@Override
		public MapCodec<TestEnvironmentDefinition.TimeOfDay> codec() {
			return CODEC;
		}
	}

	public record Weather(Type weather) implements TestEnvironmentDefinition {
		public static final MapCodec<TestEnvironmentDefinition.Weather> CODEC = RecordCodecBuilder.mapCodec(
			instance -> instance.group(Type.CODEC.fieldOf("weather").forGetter(TestEnvironmentDefinition.Weather::weather))
				.apply(instance, TestEnvironmentDefinition.Weather::new)
		);

		@Override
		public void setup(ServerLevel level) {
			this.weather.apply(level);
		}

		@Override
		public void teardown(ServerLevel level) {
			level.resetWeatherCycle();
		}

		@Override
		public MapCodec<TestEnvironmentDefinition.Weather> codec() {
			return CODEC;
		}
	}
}