tbd-station-14/Content.IntegrationTests/Tests/Interaction/InteractionTest.EntitySpecifier.cs

#nullable enable
using Content.Shared.Stacks;
using Robust.Shared.GameObjects;
using Robust.Shared.Map;
using Robust.Shared.Prototypes;

namespace Content.IntegrationTests.Tests.Interaction;

public abstract partial class InteractionTest
{
    /// <summary>
    /// Utility class for working with prototypes ids that may refer to stacks or entities.
    /// </summary>
    /// <remarks>
    /// Intended to make tests easier by removing ambiguity around "SheetSteel1", "SheetSteel", and "Steel". All three
    /// should be treated identically by interaction tests.
    /// </remarks>
    protected sealed class EntitySpecifier
    {
        /// <summary>
        /// Either the stack or entity prototype for this entity. Stack prototypes take priority.
        /// </summary>
        public string Prototype;

        /// <summary>
        /// The quantity. If the entity has a stack component, this is the total stack quantity.
        /// Otherwise this is the number of entities.
        /// </summary>
        /// <remarks>
        /// If used for spawning and this number is larger than the max stack size, only a single stack will be spawned.
        /// </remarks>
        public int Quantity;

        /// <summary>
        /// If true, a check has been performed to see if the prototype ia an entity prototype with a stack component,
        /// in which case the specifier was converted into a stack-specifier
        /// </summary>
        public bool Converted;

        public EntitySpecifier(string prototype, int quantity, bool converted = false)
        {
            Assert.That(quantity, Is.GreaterThan(0));
            Prototype = prototype;
            Quantity = quantity;
            Converted = converted;
        }

        public static implicit operator EntitySpecifier(string prototype)
            => new(prototype, 1);

        public static implicit operator EntitySpecifier((string, int) tuple)
            => new(tuple.Item1, tuple.Item2);

        /// <summary>
        /// Convert applicable entity prototypes into stack prototypes.
        /// </summary>
        public void ConvertToStack(IPrototypeManager protoMan, IComponentFactory factory)
        {
            if (Converted)
                return;

            Converted = true;

            if (string.IsNullOrWhiteSpace(Prototype))
                return;

            if (protoMan.HasIndex<StackPrototype>(Prototype))
                return;

            if (!protoMan.TryIndex<EntityPrototype>(Prototype, out var entProto))
            {
                Assert.Fail($"Unknown prototype: {Prototype}");
                return;
            }

            if (entProto.TryGetComponent<StackComponent>(factory.GetComponentName(typeof(StackComponent)),
                    out var stackComp))
            {
                Prototype = stackComp.StackTypeId;
            }
        }
    }

    protected async Task<EntityUid> SpawnEntity(EntitySpecifier spec, EntityCoordinates coords)
    {
        EntityUid uid = default!;
        if (ProtoMan.TryIndex<StackPrototype>(spec.Prototype, out var stackProto))
        {
            await Server.WaitPost(() =>
            {
                uid = SEntMan.SpawnEntity(stackProto.Spawn, coords);
                Stack.SetCount(uid, spec.Quantity);
            });
            return uid;
        }

        if (!ProtoMan.TryIndex<EntityPrototype>(spec.Prototype, out var entProto))
        {
            Assert.Fail($"Unkown prototype: {spec.Prototype}");
            return default;
        }

        if (entProto.TryGetComponent<StackComponent>(Factory.GetComponentName(typeof(StackComponent)),
                out var stackComp))
        {
            return await SpawnEntity((stackComp.StackTypeId, spec.Quantity), coords);
        }

        Assert.That(spec.Quantity, Is.EqualTo(1), "SpawnEntity only supports returning a singular entity");
        await Server.WaitPost(() => uid = SEntMan.SpawnEntity(spec.Prototype, coords));
        return uid;
    }

    /// <summary>
    /// Convert an entity-uid to a matching entity specifier. Useful when doing entity lookups & checking that the
    /// right quantity of entities/materials werre produced. Returns null if passed an entity with a null prototype.
    /// </summary>
    protected EntitySpecifier? ToEntitySpecifier(EntityUid uid)
    {
        if (SEntMan.TryGetComponent(uid, out StackComponent? stack))
            return new EntitySpecifier(stack.StackTypeId, stack.Count) { Converted = true };

        var meta = SEntMan.GetComponent<MetaDataComponent>(uid);

        if (meta.EntityPrototype is null)
            return null;

        return new(meta.EntityPrototype.ID, 1) { Converted = true };
    }
}