Project-M/Assets/_Project/Tests/EditMode/EquipSystemTests.cs

using NUnit.Framework;
using ProjectM.Server;
using ProjectM.Simulation;
using Unity.Collections;
using Unity.Core;
using Unity.Entities;
using Unity.NetCode;

namespace ProjectM.Tests
{
    /// <summary>
    /// Plain-Entities EditMode tests for the server-only <see cref="EquipSystem"/>. Seeds a player
    /// (GhostOwner + PlayerTag + InventorySlot + EquipmentSlot[4 rows] + StatModifier + AbilityRef +
    /// DefaultAbility), an inline-built ItemDatabase singleton, a mock connection, and an Equip/Unequip RPC.
    /// Pins: weapon-equip sets AbilityRef + adds the slot-tagged mod + moves the item bag->slot; unequip reverses
    /// and restores DefaultAbility; equip-over-occupied swaps the old item back; a full-bag swap is rejected with
    /// no item loss; non-equippable / absent / unresolvable-connection requests no-op (request still consumed);
    /// the unequip strip removes ONLY the slot sentinel, leaving foreign-SourceId mods (pickup 0u, upgrade) intact.
    /// </summary>
    public class EquipSystemTests
    {
        const ushort WeaponA = 100, WeaponB = 101, GearArmor = 110, Ore = 2;
        BlobAssetReference<ItemDatabaseBlob> _blob;

        [TearDown]
        public void TearDown()
        {
            if (_blob.IsCreated) _blob.Dispose();
            _blob = default;
        }

        static ItemModSpec NoMod() => new ItemModSpec { Target = 255 };

        static ItemDefBlob Mk(ushort id, byte slot, byte ability, ItemModSpec m0)
        {
            int stackMax = slot <= EquipSlotId.Tool ? 1 : 999;
            return new ItemDefBlob
            {
                ItemId = id, Category = 0, Tier = 0, StackMax = stackMax,
                EquipSlot = slot, GrantedAbilityId = ability,
                Mod0 = m0, Mod1 = NoMod(), Mod2 = NoMod(), Mod3 = NoMod(),
            };
        }

        (World world, SimulationSystemGroup group) MakeWorld(string name)
        {
            var world = new World(name);
            var group = world.GetOrCreateSystemManaged<SimulationSystemGroup>();
            group.AddSystemToUpdateList(world.GetOrCreateSystem<EquipSystem>());
            group.SortSystems();
            world.SetTime(new TimeData(elapsedTime: 0f, deltaTime: 1f / 60f));
            var em = world.EntityManager;

            var builder = new BlobBuilder(Allocator.Temp);
            ref var root = ref builder.ConstructRoot<ItemDatabaseBlob>();
            var arr = builder.Allocate(ref root.Items, 4);
            arr[0] = Mk(WeaponA, EquipSlotId.Weapon, (byte)AbilityId.FastLight, new ItemModSpec { Target = (byte)StatTarget.Damage, Op = (byte)ModOp.Flat, Value = 5f });
            arr[1] = Mk(WeaponB, EquipSlotId.Weapon, (byte)AbilityId.SlowHeavy, new ItemModSpec { Target = (byte)StatTarget.Damage, Op = (byte)ModOp.Flat, Value = 9f });
            arr[2] = Mk(GearArmor, EquipSlotId.Armor, 0, new ItemModSpec { Target = (byte)StatTarget.MoveSpeed, Op = (byte)ModOp.PercentAdd, Value = 0.1f });
            arr[3] = Mk(Ore, EquipSlotId.None, 0, NoMod());
            _blob = builder.CreateBlobAssetReference<ItemDatabaseBlob>(Allocator.Persistent);
            builder.Dispose();
            var dbE = em.CreateEntity(typeof(ItemDatabase));
            em.SetComponentData(dbE, new ItemDatabase { Value = _blob });

            return (world, group);
        }

        static Entity MakeConnection(EntityManager em, int networkId)
        {
            var e = em.CreateEntity();
            em.AddComponentData(e, new NetworkId { Value = networkId });
            return e;
        }

        static Entity MakePlayer(EntityManager em, int networkId, params (ushort id, int count)[] bagItems)
        {
            var e = em.CreateEntity();
            em.AddComponentData(e, new GhostOwner { NetworkId = networkId });
            em.AddComponent<PlayerTag>(e);
            em.AddComponentData(e, new AbilityRef { Id = (byte)AbilityId.Primary });
            em.AddComponentData(e, new DefaultAbility { Id = (byte)AbilityId.Primary });
            var bag = em.AddBuffer<InventorySlot>(e);
            foreach (var it in bagItems) bag.Add(new InventorySlot { ItemId = it.id, Count = it.count });
            var slots = em.AddBuffer<EquipmentSlot>(e);
            for (int s = 0; s < EquipSlotId.Count; s++) slots.Add(new EquipmentSlot { ItemId = 0 });
            em.AddBuffer<StatModifier>(e);
            return e;
        }

        static void MakeEquip(EntityManager em, ushort itemId, Entity conn)
        {
            var e = em.CreateEntity();
            em.AddComponentData(e, new EquipRequest { ItemId = itemId });
            em.AddComponentData(e, new ReceiveRpcCommandRequest { SourceConnection = conn });
        }

        static void MakeUnequip(EntityManager em, byte slot, Entity conn)
        {
            var e = em.CreateEntity();
            em.AddComponentData(e, new UnequipRequest { Slot = slot });
            em.AddComponentData(e, new ReceiveRpcCommandRequest { SourceConnection = conn });
        }

        static byte Ability(EntityManager em, Entity p) => em.GetComponentData<AbilityRef>(p).Id;
        static ushort Slot(EntityManager em, Entity p, byte slot) => em.GetBuffer<EquipmentSlot>(p)[slot].ItemId;
        static int Bag(EntityManager em, Entity p, ushort id) => InventoryMath.CountOf(em.GetBuffer<InventorySlot>(p), id);
        static int RequestsLeft(EntityManager em) { using var q = em.CreateEntityQuery(typeof(ReceiveRpcCommandRequest)); return q.CalculateEntityCount(); }

        static int SlotModCount(EntityManager em, Entity p, byte slot)
        {
            var mods = em.GetBuffer<StatModifier>(p);
            uint sid = Tuning.EquipSourceIdBase + (uint)slot;
            int c = 0;
            for (int i = 0; i < mods.Length; i++) if (mods[i].SourceId == sid) c++;
            return c;
        }

        static int ModCountBySource(EntityManager em, Entity p, uint sourceId)
        {
            var mods = em.GetBuffer<StatModifier>(p);
            int c = 0;
            for (int i = 0; i < mods.Length; i++) if (mods[i].SourceId == sourceId) c++;
            return c;
        }

        [Test]
        public void Equip_Weapon_Sets_Ability_Adds_Mod_Moves_Item()
        {
            var (world, group) = MakeWorld("EquipWeapon");
            using (world)
            {
                var em = world.EntityManager;
                var conn = MakeConnection(em, 1);
                var player = MakePlayer(em, 1, (WeaponA, 1));
                MakeEquip(em, WeaponA, conn);

                group.Update();

                Assert.AreEqual((byte)AbilityId.FastLight, Ability(em, player), "The weapon grants its ability into AbilityRef.");
                Assert.AreEqual(WeaponA, Slot(em, player, EquipSlotId.Weapon), "The weapon occupies the Weapon slot.");
                Assert.AreEqual(0, Bag(em, player, WeaponA), "The weapon left the bag.");
                Assert.AreEqual(1, SlotModCount(em, player, EquipSlotId.Weapon), "The weapon's mod is tagged the weapon-slot sentinel.");
                Assert.AreEqual(0, RequestsLeft(em), "The request is consumed.");
            }
        }

        [Test]
        public void Unequip_Weapon_Restores_Default_Strips_Mods_Returns_Item()
        {
            var (world, group) = MakeWorld("UnequipWeapon");
            using (world)
            {
                var em = world.EntityManager;
                var conn = MakeConnection(em, 1);
                var player = MakePlayer(em, 1, (WeaponA, 1));
                MakeEquip(em, WeaponA, conn);
                group.Update();

                MakeUnequip(em, EquipSlotId.Weapon, conn);
                group.Update();

                Assert.AreEqual((byte)AbilityId.Primary, Ability(em, player), "Unequip restores the DefaultAbility.");
                Assert.AreEqual(0, Slot(em, player, EquipSlotId.Weapon), "The Weapon slot is empty.");
                Assert.AreEqual(1, Bag(em, player, WeaponA), "The weapon is back in the bag.");
                Assert.AreEqual(0, SlotModCount(em, player, EquipSlotId.Weapon), "The weapon's mod is stripped.");
            }
        }

        [Test]
        public void Equip_Over_Occupied_Swaps_Old_Item_Back()
        {
            var (world, group) = MakeWorld("EquipSwap");
            using (world)
            {
                var em = world.EntityManager;
                var conn = MakeConnection(em, 1);
                var player = MakePlayer(em, 1, (WeaponA, 1), (WeaponB, 1));
                MakeEquip(em, WeaponA, conn);
                group.Update();

                MakeEquip(em, WeaponB, conn);
                group.Update();

                Assert.AreEqual(WeaponB, Slot(em, player, EquipSlotId.Weapon), "Weapon B now occupies the slot.");
                Assert.AreEqual((byte)AbilityId.SlowHeavy, Ability(em, player), "AbilityRef swaps to weapon B's ability.");
                Assert.AreEqual(1, Bag(em, player, WeaponA), "Weapon A swapped back into the bag.");
                Assert.AreEqual(0, Bag(em, player, WeaponB), "Weapon B left the bag.");
                Assert.AreEqual(1, SlotModCount(em, player, EquipSlotId.Weapon), "Exactly weapon B's single mod remains (A's stripped).");
            }
        }

        [Test]
        public void Swap_With_Full_Bag_Is_Rejected_No_Item_Loss()
        {
            var (world, group) = MakeWorld("FullBagSwap");
            using (world)
            {
                var em = world.EntityManager;
                var conn = MakeConnection(em, 1);
                // Pre-equip weapon A directly, then fill the bag completely (incl. weapon B). Equipping B must
                // reject because the bag has no room to receive the swapped-out weapon A.
                var player = MakePlayer(em, 1, (WeaponB, 1));
                var preSlots = em.GetBuffer<EquipmentSlot>(player);
                preSlots[EquipSlotId.Weapon] = new EquipmentSlot { ItemId = WeaponA };
                var bag = em.GetBuffer<InventorySlot>(player);
                for (int i = 0; bag.Length < Tuning.InventoryMaxSlots; i++)
                    bag.Add(new InventorySlot { ItemId = (ushort)(200 + i), Count = 1 });

                MakeEquip(em, WeaponB, conn);
                group.Update();

                Assert.AreEqual(WeaponA, Slot(em, player, EquipSlotId.Weapon), "The occupied slot is unchanged (equip rejected).");
                Assert.AreEqual(1, Bag(em, player, WeaponB), "Weapon B was NOT withdrawn — no item loss.");
                Assert.AreEqual(0, RequestsLeft(em), "The request is still consumed.");
            }
        }

        [Test]
        public void Equip_NonEquippable_Or_Absent_Item_Is_NoOp()
        {
            var (world, group) = MakeWorld("NoOpEquip");
            using (world)
            {
                var em = world.EntityManager;
                var conn = MakeConnection(em, 1);
                var player = MakePlayer(em, 1, (Ore, 5)); // carries a resource (EquipSlot=None) but no weapon
                MakeEquip(em, Ore, conn);       // not equippable
                MakeEquip(em, WeaponA, conn);   // not in the bag
                group.Update();

                Assert.AreEqual(0, Slot(em, player, EquipSlotId.Weapon), "Nothing equipped.");
                Assert.AreEqual(5, Bag(em, player, Ore), "The resource is untouched.");
                Assert.AreEqual((byte)AbilityId.Primary, Ability(em, player), "AbilityRef unchanged.");
                Assert.AreEqual(0, RequestsLeft(em), "Both requests are consumed.");
            }
        }

        [Test]
        public void Equip_From_Unresolvable_Connection_NoOp()
        {
            var (world, group) = MakeWorld("UnresolvedEquip");
            using (world)
            {
                var em = world.EntityManager;
                var player = MakePlayer(em, 1, (WeaponA, 1));
                MakeEquip(em, WeaponA, Entity.Null); // no NetworkId on Entity.Null

                group.Update();

                Assert.AreEqual(0, Slot(em, player, EquipSlotId.Weapon), "An unresolvable sender equips nothing.");
                Assert.AreEqual(1, Bag(em, player, WeaponA), "The item stays in the bag.");
                Assert.AreEqual(0, RequestsLeft(em), "The request is still consumed.");
            }
        }

        [Test]
        public void Strip_Removes_Only_The_Slot_Sentinel_Leaving_Foreign_Mods()
        {
            var (world, group) = MakeWorld("StripIsolation");
            using (world)
            {
                var em = world.EntityManager;
                var conn = MakeConnection(em, 1);
                var player = MakePlayer(em, 1, (GearArmor, 1));
                // Seed foreign modifiers that unequip must NOT touch: a pickup (SourceId 0) + the ability upgrade.
                var mods = em.GetBuffer<StatModifier>(player);
                mods.Add(new StatModifier { Target = (byte)StatTarget.Damage, Op = (byte)ModOp.Flat, Value = 3f, SourceId = 0u });
                mods.Add(new StatModifier { Target = (byte)StatTarget.Damage, Op = (byte)ModOp.PercentAdd, Value = 0.25f, SourceId = Tuning.AbilityUpgradeSourceId });

                MakeEquip(em, GearArmor, conn);
                group.Update();
                Assert.AreEqual(1, SlotModCount(em, player, EquipSlotId.Armor), "Gear adds its armor-slot mod.");

                MakeUnequip(em, EquipSlotId.Armor, conn);
                group.Update();

                Assert.AreEqual(0, SlotModCount(em, player, EquipSlotId.Armor), "Unequip strips the armor-slot mod.");
                Assert.AreEqual(1, ModCountBySource(em, player, 0u), "The pickup mod (SourceId 0) is untouched.");
                Assert.AreEqual(1, ModCountBySource(em, player, Tuning.AbilityUpgradeSourceId), "The upgrade mod is untouched.");
                Assert.AreEqual(1, Bag(em, player, GearArmor), "The gear returns to the bag.");
            }
        }
    }
}