Fix lathe arbitrage test (#34449)

* Fix lathe arbitrage test

* Add refinables

* nullable

* nullable2

* Fix merge

* Ignore failures

Content.IntegrationTests/Tests/MaterialArbitrageTest.cs

@@ -1,12 +1,13 @@
+#nullable enable
 using System.Collections.Generic;
 using Content.Server.Cargo.Systems;
 using Content.Server.Construction.Completions;
 using Content.Server.Construction.Components;
 using Content.Server.Destructible;
 using Content.Server.Destructible.Thresholds.Behaviors;
+using Content.Server.Lathe;
 using Content.Server.Stack;
 using Content.Shared.Chemistry.Reagent;
-using Content.Shared.Construction.Components;
 using Content.Shared.Construction.Prototypes;
 using Content.Shared.Construction.Steps;
 using Content.Shared.FixedPoint;
@@ -14,10 +15,9 @@ using Content.Shared.Lathe;
 using Content.Shared.Materials;
 using Content.Shared.Research.Prototypes;
 using Content.Shared.Stacks;
+using Content.Shared.Tools.Components;
 using Robust.Shared.GameObjects;
-using Robust.Shared.Map;
 using Robust.Shared.Prototypes;
-using Robust.Shared.Utility;
 
 namespace Content.IntegrationTests.Tests;
 
@@ -28,6 +28,21 @@ namespace Content.IntegrationTests.Tests;
 [TestFixture]
 public sealed class MaterialArbitrageTest
 {
+    // These recipes are currently broken and need fixing. You should not be adding to these sets.
+    private readonly HashSet<string> _destructionArbitrageIgnore =
+    [
+        "BaseChemistryEmptyVial", "DrinkShotGlass", "Beaker", "SodiumLightTube", "DrinkGlassCoupeShaped",
+        "LedLightBulb", "ExteriorLightTube", "LightTube", "DrinkGlass", "DimLightBulb", "LightBulb", "LedLightTube",
+        "SheetRGlass1", "ChemistryEmptyBottle01", "WarmLightBulb",
+    ];
+
+    private readonly HashSet<string> _compositionArbitrageIgnore =
+    [
+        "FoodPlateSmall", "AirTank", "FoodPlateTin", "FoodPlateMuffinTin", "WeaponCapacitorRechargerCircuitboard",
+        "WeaponCapacitorRechargerCircuitboard", "BorgChargerCircuitboard", "BorgChargerCircuitboard", "FoodPlate",
+        "CellRechargerCircuitboard", "CellRechargerCircuitboard",
+    ];
+
     [Test]
     public async Task NoMaterialArbitrage()
     {
@@ -38,13 +53,12 @@ public sealed class MaterialArbitrageTest
         await server.WaitIdleAsync();
 
         var entManager = server.ResolveDependency<IEntityManager>();
-        var mapManager = server.ResolveDependency<IMapManager>();
         var protoManager = server.ResolveDependency<IPrototypeManager>();
 
         var pricing = entManager.System<PricingSystem>();
         var stackSys = entManager.System<StackSystem>();
         var mapSystem = server.System<SharedMapSystem>();
-        var latheSys = server.System<SharedLatheSystem>();
+        var latheSys = server.System<LatheSystem>();
         var compFact = server.ResolveDependency<IComponentFactory>();
 
         Assert.That(mapSystem.IsInitialized(testMap.MapId));
@@ -53,13 +67,24 @@ public sealed class MaterialArbitrageTest
         var compositionName = compFact.GetComponentName(typeof(PhysicalCompositionComponent));
         var materialName = compFact.GetComponentName(typeof(MaterialComponent));
         var destructibleName = compFact.GetComponentName(typeof(DestructibleComponent));
+        var refinableName = compFact.GetComponentName(typeof(ToolRefinableComponent));
 
         // get the inverted lathe recipe dictionary
         var latheRecipes = latheSys.InverseRecipes;
 
-        // Lets assume the possible lathe for resource multipliers:
+        // Find the lowest multiplier / optimal lathe that can be used to construct a recipie.
-        // TODO: each recipe can technically have its own cost multiplier associated with it, so this test needs redone to factor that in.
+        var minMultiplier = new Dictionary<ProtoId<LatheRecipePrototype>, float>();
-        var multiplier = MathF.Pow(0.85f, 3);
+
+        foreach (var (_, lathe) in pair.GetPrototypesWithComponent<LatheComponent>())
+        {
+            foreach (var recipe in latheSys.GetAllPossibleRecipes(lathe))
+            {
+                if (!minMultiplier.TryGetValue(recipe, out var min))
+                    min = 1;
+
+                minMultiplier[recipe] = Math.Min(min, lathe.MaterialUseMultiplier);
+            }
+        }
 
         // create construction dictionary
         Dictionary<string, ConstructionComponent> constructionRecipes = new();
@@ -122,6 +147,65 @@ public sealed class MaterialArbitrageTest
 
         Dictionary<string, (Dictionary<string, int> Ents, Dictionary<string, int> Mats)> spawnedOnDestroy = new();
 
+        // cache the compositions of entities
+        // If the entity is refineable (i.e. glass shared can be turned into glass, we take the greater of the two compositions.
+        Dictionary<EntProtoId, Dictionary<string, int>> compositions = new();
+        foreach (var proto in protoManager.EnumeratePrototypes<EntityPrototype>())
+        {
+            Dictionary<string, int>? baseComposition = null;
+
+            if (proto.Components.ContainsKey(materialName)
+                && proto.Components.TryGetValue(compositionName, out var compositionReg))
+            {
+                var compositionComp = (PhysicalCompositionComponent)compositionReg.Component;
+                baseComposition = compositionComp.MaterialComposition;
+
+            }
+
+            if (!proto.Components.TryGetValue(refinableName, out var refinableReg))
+            {
+                if (baseComposition != null)
+                    compositions[proto.ID] = new(baseComposition);
+                continue;
+            }
+
+            var composition = new Dictionary<string, int>();
+            compositions.Add(proto.ID, composition);
+
+            var refinable = (ToolRefinableComponent)refinableReg.Component;
+            foreach (var refineResult in refinable.RefineResult)
+            {
+                if (refineResult.PrototypeId == null)
+                    continue;
+
+                var refineProto = protoManager.Index(refineResult.PrototypeId.Value);
+                if (!refineProto.Components.ContainsKey(materialName))
+                    continue;
+
+                if (!refineProto.Components.TryGetValue(compositionName, out var refinedCompositionReg))
+                    continue;
+
+                var refinedCompositionComp = (PhysicalCompositionComponent)refinedCompositionReg.Component;
+
+                // This assumes refine results do not have complex spawn behaviours like exclusive groups.
+                var quantity = refineResult.MaxAmount;
+
+                foreach (var (matId, amount) in refinedCompositionComp.MaterialComposition)
+                {
+                    composition[matId] = quantity * amount + composition.GetValueOrDefault(matId);
+                }
+            }
+
+            if (baseComposition == null)
+                continue;
+
+            // If the un-refined material quantity is greater than the refined quantity, we use that instead.
+            foreach (var (matId, amount) in baseComposition)
+            {
+                composition[matId] = Math.Max(amount, composition.GetValueOrDefault(matId));
+            }
+        }
+
         // Here we get the set of entities/materials spawned when destroying an entity.
         foreach (var proto in protoManager.EnumeratePrototypes<EntityPrototype>())
         {
@@ -151,16 +235,10 @@ public sealed class MaterialArbitrageTest
                     {
                         spawnedEnts[key] = spawnedEnts.GetValueOrDefault(key) + value.Max;
 
-                        var spawnProto = protoManager.Index<EntityPrototype>(key);
+                        if (!compositions.TryGetValue(key, out var composition))
-
-                        // get the amount of each material included in the entity
-
-                        if (!spawnProto.Components.ContainsKey(materialName) ||
-                            !spawnProto.Components.TryGetValue(compositionName, out var compositionReg))
                             continue;
 
-                        var mat = (PhysicalCompositionComponent) compositionReg.Component;
+                        foreach (var (matId, amount) in composition)
-                        foreach (var (matId, amount) in mat.MaterialComposition)
                         {
                             spawnedMats[matId] = value.Max * amount + spawnedMats.GetValueOrDefault(matId);
                         }
@@ -173,10 +251,13 @@ public sealed class MaterialArbitrageTest
         }
 
         // This is the main loop where we actually check for destruction arbitrage
-        Assert.Multiple(async () =>
+        await Assert.MultipleAsync(async () =>
         {
             foreach (var (id, (spawnedEnts, spawnedMats)) in spawnedOnDestroy)
             {
+                if (_destructionArbitrageIgnore.Contains(id))
+                    continue;
+
                 // Check cargo sell price
                 // several constructible entities have no sell price
                 // also this test only really matters if the entity is also purchaseable.... eh..
@@ -190,6 +271,11 @@ public sealed class MaterialArbitrageTest
                 {
                     foreach (var recipe in recipes)
                     {
+                        if (!minMultiplier.TryGetValue(recipe, out var multiplier))
+                        {
+                            server.Log.Info($"Unused lathe recipe? {recipe.ID}?");
+                            continue;
+                        }
                         foreach (var (matId, amount) in recipe.Materials)
                         {
                             var actualAmount = SharedLatheSystem.AdjustMaterial(amount, recipe.ApplyMaterialDiscount, multiplier);
@@ -231,6 +317,9 @@ public sealed class MaterialArbitrageTest
                 var edge = cur.GetEdge(node.Name);
                 cur = node;
 
+                if (edge == null)
+                    continue;
+
                 foreach (var completion in edge.Completed)
                 {
                     if (completion is not SpawnPrototype spawnCompletion)
@@ -253,9 +342,9 @@ public sealed class MaterialArbitrageTest
         }
 
         // This is functionally the same loop as before, but now testing deconstruction rather than destruction.
-        // This is pretty braindead. In principle construction graphs can have loops and whatnot.
+        // This is pretty brain-dead. In principle construction graphs can have loops and whatnot.
 
-        Assert.Multiple(async () =>
+        await Assert.MultipleAsync(async () =>
         {
             foreach (var (id, deconstructedMats) in deconstructionMaterials)
             {
@@ -270,6 +359,11 @@ public sealed class MaterialArbitrageTest
                 {
                     foreach (var recipe in recipes)
                     {
+                        if (!minMultiplier.TryGetValue(recipe, out var multiplier))
+                        {
+                            server.Log.Info($"Unused lathe recipe? {recipe.ID}?");
+                            continue;
+                        }
                         foreach (var (matId, amount) in recipe.Materials)
                         {
                             var actualAmount = SharedLatheSystem.AdjustMaterial(amount, recipe.ApplyMaterialDiscount, multiplier);
@@ -291,7 +385,7 @@ public sealed class MaterialArbitrageTest
             }
         });
 
-        // create phyiscal composition dictionary
+        // create physical composition dictionary
         // this doesn't account for the chemicals in the composition
         Dictionary<string, PhysicalCompositionComponent> physicalCompositions = new();
         foreach (var proto in protoManager.EnumeratePrototypes<EntityPrototype>())
@@ -308,10 +402,13 @@ public sealed class MaterialArbitrageTest
 
         // This is functionally the same loop as before, but now testing composition rather than destruction or deconstruction.
         // This doesn't take into account chemicals generated when deconstructing. Maybe it should.
-        Assert.Multiple(async () =>
+        await Assert.MultipleAsync(async () =>
         {
             foreach (var (id, compositionComponent) in physicalCompositions)
             {
+                if (_compositionArbitrageIgnore.Contains(id))
+                    continue;
+
                 // Check cargo sell price
                 var materialPrice = await GetDeconstructedPrice(compositionComponent.MaterialComposition);
                 var chemicalPrice = await GetChemicalCompositionPrice(compositionComponent.ChemicalComposition);
@@ -325,6 +422,11 @@ public sealed class MaterialArbitrageTest
                 {
                     foreach (var recipe in recipes)
                     {
+                        if (!minMultiplier.TryGetValue(recipe, out var multiplier))
+                        {
+                            server.Log.Info($"Unused lathe recipe? {recipe.ID}?");
+                            continue;
+                        }
                         foreach (var (matId, amount) in recipe.Materials)
                         {
                             var actualAmount = SharedLatheSystem.AdjustMaterial(amount, recipe.ApplyMaterialDiscount, multiplier);

Content.Server/Lathe/LatheSystem.cs

@@ -61,7 +61,6 @@ namespace Content.Server.Lathe
         /// Per-tick cache
         /// </summary>
         private readonly List<GasMixture> _environments = new();
-        private readonly HashSet<ProtoId<LatheRecipePrototype>> _availableRecipes = new();
 
         public override void Initialize()
         {
@@ -162,12 +161,8 @@ namespace Content.Server.Lathe
 
         public List<ProtoId<LatheRecipePrototype>> GetAvailableRecipes(EntityUid uid, LatheComponent component, bool getUnavailable = false)
         {
-            _availableRecipes.Clear();
+            var ev = new LatheGetRecipesEvent((uid, component), getUnavailable);
-            AddRecipesFromPacks(_availableRecipes, component.StaticPacks);
+            AddRecipesFromPacks(ev.Recipes, component.StaticPacks);
-            var ev = new LatheGetRecipesEvent(uid, getUnavailable)
-            {
-                Recipes = _availableRecipes
-            };
             RaiseLocalEvent(uid, ev);
             return ev.Recipes.ToList();
         }
@@ -290,7 +285,7 @@ namespace Content.Server.Lathe
                 var pack = _proto.Index(id);
                 foreach (var recipe in pack.Recipes)
                 {
-                    if (args.getUnavailable || database.UnlockedRecipes.Contains(recipe))
+                    if (args.GetUnavailable || database.UnlockedRecipes.Contains(recipe))
                         args.Recipes.Add(recipe);
                 }
             }
@@ -298,10 +293,8 @@ namespace Content.Server.Lathe
 
         private void OnGetRecipes(EntityUid uid, TechnologyDatabaseComponent component, LatheGetRecipesEvent args)
         {
-            if (uid != args.Lathe || !TryComp<LatheComponent>(uid, out var latheComponent))
+            if (uid == args.Lathe)
-                return;
+                AddRecipesFromDynamicPacks(ref args, component, args.Comp.DynamicPacks);
-
-            AddRecipesFromDynamicPacks(ref args, component, latheComponent.DynamicPacks);
         }
 
         private void GetEmagLatheRecipes(EntityUid uid, EmagLatheRecipesComponent component, LatheGetRecipesEvent args)
@@ -309,7 +302,7 @@ namespace Content.Server.Lathe
             if (uid != args.Lathe)
                 return;
 
-            if (!args.getUnavailable && !_emag.CheckFlag(uid, EmagType.Interaction))
+            if (!args.GetUnavailable && !_emag.CheckFlag(uid, EmagType.Interaction))
                 return;
 
             AddRecipesFromPacks(args.Recipes, component.EmagStaticPacks);

Content.Shared/Lathe/LatheComponent.cs

@@ -21,6 +21,9 @@ namespace Content.Shared.Lathe
         /// </summary>
         [DataField]
         public List<ProtoId<LatheRecipePackPrototype>> DynamicPacks = new();
+        // Note that this shouldn't be modified dynamically.
+        // I.e., this + the static recipies should represent all recipies that the lathe can ever make
+        // Otherwise the material arbitrage test and/or LatheSystem.GetAllBaseRecipes needs to be updated
 
         /// <summary>
         /// The lathe's construction queue
@@ -81,15 +84,16 @@ namespace Content.Shared.Lathe
     public sealed class LatheGetRecipesEvent : EntityEventArgs
     {
         public readonly EntityUid Lathe;
+        public readonly LatheComponent Comp;
 
-        public bool getUnavailable;
+        public bool GetUnavailable;
 
         public HashSet<ProtoId<LatheRecipePrototype>> Recipes = new();
 
-        public LatheGetRecipesEvent(EntityUid lathe, bool forced)
+        public LatheGetRecipesEvent(Entity<LatheComponent> lathe, bool forced)
         {
-            Lathe = lathe;
+            (Lathe, Comp) = lathe;
-            getUnavailable = forced;
+            GetUnavailable = forced;
         }
     }
 

Content.Shared/Lathe/SharedLatheSystem.cs

@@ -34,6 +34,25 @@ public abstract class SharedLatheSystem : EntitySystem
         BuildInverseRecipeDictionary();
     }
 
+    /// <summary>
+    /// Get the set of all recipes that a lathe could possibly ever create (e.g., if all techs were unlocked).
+    /// </summary>
+    public HashSet<ProtoId<LatheRecipePrototype>> GetAllPossibleRecipes(LatheComponent component)
+    {
+        var recipes = new HashSet<ProtoId<LatheRecipePrototype>>();
+        foreach (var pack in component.StaticPacks)
+        {
+            recipes.UnionWith(_proto.Index(pack).Recipes);
+        }
+
+        foreach (var pack in component.DynamicPacks)
+        {
+            recipes.UnionWith(_proto.Index(pack).Recipes);
+        }
+
+        return recipes;
+    }
+
     /// <summary>
     /// Add every recipe in the list of recipe packs to a single hashset.
     /// </summary>