CatStation/Content.IntegrationTests/Tests/VendingMachineRestockTest.cs

#nullable enable
using System.Collections.Generic;
using Content.IntegrationTests.Fixtures;
using Content.Server.VendingMachines;
using Content.Server.Wires;
using Content.Shared.Cargo.Prototypes;
using Content.Shared.Containers;
using Content.Shared.Damage;
using Content.Shared.Damage.Prototypes;
using Content.Shared.Damage.Systems;
using Content.Shared.EntityTable;
using Content.Shared.Prototypes;
using Content.Shared.Storage.EntitySystems;
using Content.Shared.VendingMachines;
using Content.Shared.VendingMachines.Components;
using Content.Shared.Wires;
using Robust.Shared.GameObjects;
using Robust.Shared.Prototypes;

namespace Content.IntegrationTests.Tests
{
    [TestFixture]
    [TestOf(typeof(VendingMachineRestockComponent))]
    [TestOf(typeof(VendingMachineSystem))]
    public sealed class VendingMachineRestockTest : GameTest
    {
        private static readonly ProtoId<DamageTypePrototype> TestDamageType = "Blunt";

        [TestPrototypes]
        private const string Prototypes = @"
- type: entity
  name: HumanVendingDummy
  id: HumanVendingDummy
  components:
  - type: Hands
  - type: ComplexInteraction
  - type: Body
    prototype: Human

- type: entity
  parent: FoodSnackBase
  id: TestRamen
  name: TestRamen

- type: vendingMachineInventory
  id: TestInventory
  startingInventory:
    TestRamen: 1

- type: vendingMachineInventory
  id: OtherTestInventory
  startingInventory:
    TestRamen: 3

- type: vendingMachineInventory
  id: BigTestInventory
  startingInventory:
    TestRamen: 4

- type: entity
  parent: BaseVendingMachineRestock
  id: TestRestockWrong
  name: TestRestockWrong
  components:
  - type: VendingMachineRestock
    canRestock:
    - OtherTestInventory

- type: entity
  parent: BaseVendingMachineRestock
  id: TestRestockCorrect
  name: TestRestockCorrect
  components:
  - type: VendingMachineRestock
    canRestock:
    - TestInventory

- type: entity
  parent: BaseVendingMachineRestock
  id: TestRestockExplode
  name: TestRestockExplode
  components:
  - type: Damageable
    damageModifierSet: Metallic
  - type: Injurable
    damageContainer: Inorganic
  - type: Destructible
    thresholds:
    - trigger:
        !type:DamageTrigger
        damage: 20
      behaviors:
      - !type:DumpRestockInventory
      - !type:DoActsBehavior
        acts: [ 'Destruction' ]
  - type: VendingMachineRestock
    canRestock:
    - BigTestInventory

- type: entity
  parent: BaseVendingMachine
  id: VendingMachineTest
  name: Test Ramen
  components:
  - type: Wires
    layoutId: Vending
  - type: VendingMachine
    pack: TestInventory
  - type: VendingMachineEject
  - type: Sprite
    sprite: error.rsi
";

        [Test]
        public async Task TestAllRestocksAreAvailableToBuy()
        {
            var pair = Pair;
            var server = pair.Server;
            await server.WaitIdleAsync();

            var prototypeManager = server.ResolveDependency<IPrototypeManager>();
            var compFact = server.ResolveDependency<IComponentFactory>();
            var entityTable = server.EntMan.System<EntityTableSystem>();

            await server.WaitAssertion(() =>
            {
                // Collect all entity prototypes which are vending machine restocks.
                var restockEntities = new HashSet<EntProtoId<VendingMachineRestockComponent>>();
                foreach (var proto in prototypeManager.EnumeratePrototypes<EntityPrototype>())
                {
                    if (proto.Abstract
                        || pair.IsTestPrototype(proto)
                        || !proto.HasComponent<VendingMachineRestockComponent>())
                        continue;

                    restockEntities.Add(proto.ID);
                }

                // Collect all entity prototypes with `EntityTableContainerFill`s which contain those restock entities.
                // Specifically, this is a mapping of entities-with-container-fill to their-contained-entities-which-are-restocks.
                Dictionary<EntProtoId<EntityTableContainerFillComponent>,
                    List<EntProtoId<VendingMachineRestockComponent>>> entitiesWhichSpawnRestocks = new();
                foreach (var proto in prototypeManager.EnumeratePrototypes<EntityPrototype>())
                {
                    if (!proto.TryComp<EntityTableContainerFillComponent>(out var fill, compFact))
                        continue;

                    var containers = fill.Containers;

                    // We only care about the special known container.
                    if (!containers.TryGetValue(SharedEntityStorageSystem.ContainerName, out var container))
                        continue;

                    var entitiesInProtoContainingRestock = new List<EntProtoId<VendingMachineRestockComponent>>();
                    foreach (var (fillSpawnEntry, _) in entityTable.ListSpawns(container))
                    {
                        if (restockEntities.Contains(fillSpawnEntry.Id))
                            entitiesInProtoContainingRestock.Add(fillSpawnEntry.Id);
                    }

                    if (entitiesInProtoContainingRestock.Count > 0)
                        entitiesWhichSpawnRestocks.Add(proto.ID, entitiesInProtoContainingRestock);
                }

                // Remove all restock entities from our set which are either directly purchasable as a CargoProduct, or
                // which are spawned by EntityTableContainerFill on a CargoProduct.
                foreach (var proto in prototypeManager.EnumeratePrototypes<CargoProductPrototype>())
                {
                    // If the cargo product's product is the restock itself, just remove it.
                    restockEntities.Remove(proto.Product.Id);

                    // Check if the product is an entity which spawns a restock.
                    if (entitiesWhichSpawnRestocks.TryGetValue(proto.Product.Id, out var restocksSpawnedByProduct))
                    {
                        foreach (var entry in restocksSpawnedByProduct)
                        {
                            restockEntities.Remove(entry);
                        }

                        entitiesWhichSpawnRestocks.Remove(proto.Product.Id);
                    }
                }
                // Any entities left in restockEntities are restocks which can't be bought from Cargo.

                Assert.Multiple(() =>
                {
                    const string restockCompName = nameof(VendingMachineRestockComponent);

                    Assert.That(entitiesWhichSpawnRestocks,
                        Has.Count.EqualTo(0),
                        $"Some entities containing entities with {restockCompName} are unavailable for purchase: \n - {string.Join("\n - ", entitiesWhichSpawnRestocks.Keys)}");



                    Assert.That(restockEntities,
                        Has.Count.EqualTo(0),
                        $"Some entities with {restockCompName} are unavailable for purchase: \n - {string.Join("\n - ", restockEntities)}");
                });
            });
        }

        [Test]
        public async Task TestCompleteRestockProcess()
        {
            var pair = Pair;
            var server = pair.Server;
            await server.WaitIdleAsync();

            var entityManager = server.ResolveDependency<IEntityManager>();
            var entitySystemManager = server.ResolveDependency<IEntitySystemManager>();
            var mapSystem = server.System<SharedMapSystem>();

            EntityUid packageRight;
            EntityUid packageWrong;
            EntityUid machine;
            EntityUid user;
            VendingMachineComponent machineComponent = default!;
            VendingMachineRestockComponent restockRightComponent = default!;
            VendingMachineRestockComponent restockWrongComponent = default!;
            WiresPanelComponent machineWiresPanel = default!;

            var testMap = await pair.CreateTestMap();

            await server.WaitAssertion(() =>
            {
                var coordinates = testMap.GridCoords;

                // Spawn the entities.
                user = entityManager.SpawnEntity("HumanVendingDummy", coordinates);
                machine = entityManager.SpawnEntity("VendingMachineTest", coordinates);
                packageRight = entityManager.SpawnEntity("TestRestockCorrect", coordinates);
                packageWrong = entityManager.SpawnEntity("TestRestockWrong", coordinates);

                // Sanity test for components existing.
                Assert.Multiple(() =>
                {
                    Assert.That(entityManager.TryGetComponent(machine, out machineComponent!), $"Machine has no {nameof(VendingMachineComponent)}");
                    Assert.That(entityManager.TryGetComponent(packageRight, out restockRightComponent!), $"Correct package has no {nameof(VendingMachineRestockComponent)}");
                    Assert.That(entityManager.TryGetComponent(packageWrong, out restockWrongComponent!), $"Wrong package has no {nameof(VendingMachineRestockComponent)}");
                    Assert.That(entityManager.TryGetComponent(machine, out machineWiresPanel!), $"Machine has no {nameof(WiresPanelComponent)}");
                });

                var systemMachine = entitySystemManager.GetEntitySystem<VendingMachineSystem>();

                // Test that the panel needs to be opened first.
                Assert.Multiple(() =>
                {
                    Assert.That(systemMachine.TryAccessMachine(packageRight, restockRightComponent, machineComponent, user, machine), Is.False, "Right package is able to restock without opened access panel");
                    Assert.That(systemMachine.TryAccessMachine(packageWrong, restockWrongComponent, machineComponent, user, machine), Is.False, "Wrong package is able to restock without opened access panel");
                });

                var systemWires = entitySystemManager.GetEntitySystem<WiresSystem>();
                // Open the panel.
                systemWires.TogglePanel(machine, machineWiresPanel, true);

                Assert.Multiple(() =>
                {
                    // Test that the right package works for the right machine.
                    Assert.That(systemMachine.TryAccessMachine(packageRight, restockRightComponent, machineComponent, user, machine), Is.True, "Correct package is unable to restock with access panel opened");

                    // Test that the wrong package does not work.
                    Assert.That(systemMachine.TryMatchPackageToMachine(packageWrong, restockWrongComponent, machineComponent, user, machine), Is.False, "Package with invalid canRestock is able to restock machine");

                    // Test that the right package does work.
                    Assert.That(systemMachine.TryMatchPackageToMachine(packageRight, restockRightComponent, machineComponent, user, machine), Is.True, "Package with valid canRestock is unable to restock machine");

                    // Make sure there's something in there to begin with.
                    Assert.That(systemMachine.GetAvailableInventory(machine, machineComponent), Has.Count.GreaterThan(0),
                        "Machine inventory is empty before emptying.");
                });

                // Empty the inventory.
                systemMachine.EjectRandom((machine, machineComponent), false, true);
                Assert.That(systemMachine.GetAvailableInventory(machine, machineComponent), Has.Count.EqualTo(0),
                    "Machine inventory is not empty after ejecting.");

                // Test that the inventory is actually restocked.
                systemMachine.TryRestockInventory(machine, machineComponent);
                Assert.That(systemMachine.GetAvailableInventory(machine, machineComponent), Has.Count.GreaterThan(0),
                    "Machine available inventory count is not greater than zero after restock.");

                mapSystem.DeleteMap(testMap.MapId);
            });
        }

        [Test]
        public async Task TestRestockBreaksOpen()
        {
            var pair = Pair;
            var server = pair.Server;
            await server.WaitIdleAsync();

            var prototypeManager = server.ResolveDependency<IPrototypeManager>();
            var entityManager = server.ResolveDependency<IEntityManager>();
            var entitySystemManager = server.ResolveDependency<IEntitySystemManager>();
            var mapSystem = server.System<SharedMapSystem>();

            var damageableSystem = entitySystemManager.GetEntitySystem<DamageableSystem>();

            var testMap = await pair.CreateTestMap();

            EntityUid restock = default;

            await server.WaitAssertion(() =>
            {
                var coordinates = testMap.GridCoords;

                var totalStartingRamen = 0;

                foreach (var meta in entityManager.EntityQuery<MetaDataComponent>())
                    if (!meta.Deleted && meta.EntityPrototype?.ID == "TestRamen")
                        totalStartingRamen++;

                Assert.That(totalStartingRamen, Is.EqualTo(0),
                    "Did not start with zero ramen.");

                restock = entityManager.SpawnEntity("TestRestockExplode", coordinates);
                var damageSpec = new DamageSpecifier(prototypeManager.Index(TestDamageType), 100);
                var damageResult = damageableSystem.ChangeDamage(restock, damageSpec);

#pragma warning disable NUnit2045
                Assert.That(!damageResult.Empty, "Received empty damageResult when attempting to damage restock box.");

                Assert.That((int)damageResult.GetTotal(), Is.GreaterThan(0), "Box damage result was not greater than 0.");
#pragma warning restore NUnit2045
            });
            await server.WaitRunTicks(15);
            await server.WaitAssertion(() =>
            {
                Assert.That(entityManager.Deleted(restock),
                    "Restock box was not deleted after being damaged.");

                var totalRamen = 0;

                foreach (var meta in entityManager.EntityQuery<MetaDataComponent>())
                    if (!meta.Deleted && meta.EntityPrototype?.ID == "TestRamen")
                        totalRamen++;

                Assert.That(totalRamen, Is.EqualTo(2),
                    "Did not find enough ramen after destroying restock box.");

                mapSystem.DeleteMap(testMap.MapId);
            });
        }

        [Test]
        public async Task TestRestockInventoryBounds()
        {
            var pair = Pair;
            var server = pair.Server;
            await server.WaitIdleAsync();

            var entityManager = server.ResolveDependency<IEntityManager>();
            var entitySystemManager = server.ResolveDependency<IEntitySystemManager>();

            var vendingMachineSystem = entitySystemManager.GetEntitySystem<SharedVendingMachineSystem>();

            var testMap = await pair.CreateTestMap();

            await server.WaitAssertion(() =>
            {
                var coordinates = testMap.GridCoords;

                var machine = entityManager.SpawnEntity("VendingMachineTest", coordinates);

                Assert.That(vendingMachineSystem.GetAvailableInventory(machine), Has.Count.EqualTo(1),
                    "Machine's available inventory did not contain one entry.");

                Assert.That(vendingMachineSystem.GetAvailableInventory(machine)[0].Amount, Is.EqualTo(1),
                    "Machine's available inventory is not the expected amount.");

                vendingMachineSystem.RestockInventoryFromPrototype(machine);

                Assert.That(vendingMachineSystem.GetAvailableInventory(machine)[0].Amount, Is.EqualTo(2),
                    "Machine's available inventory is not double its starting amount after a restock.");

                vendingMachineSystem.RestockInventoryFromPrototype(machine);

                Assert.That(vendingMachineSystem.GetAvailableInventory(machine)[0].Amount, Is.EqualTo(3),
                    "Machine's available inventory is not triple its starting amount after two restocks.");

                vendingMachineSystem.RestockInventoryFromPrototype(machine);

                Assert.That(vendingMachineSystem.GetAvailableInventory(machine)[0].Amount, Is.EqualTo(3),
                    "Machine's available inventory did not stay the same after a third restock.");
            });
        }
    }
}