Fix DragonHandler ignoring multiple dragons (#798)

More fallout from #792 - I forgot that you can trade for multiple dragons at once in treasure trade. After #792 this only gives 1 dragon and scams the player, as DragonHandler ignores `Entity.Quantity`. This is actually quite tricky to handle since each dragon could put you over the limit. This change special-cases dragon trades and splits them up to be sent to the batch reward handler one-by-one which can process each one individually. To accomplish this the BatchGrantReward method has been expanded to use a fallback for non-batchable entities instead of just throwing.
SapiensAnatis · May 16, 2024 · b9d0d82 · b9d0d82
1 parent f2c7a4e
commit b9d0d82
Show file tree

Hide file tree

Showing 11 changed files with 139 additions and 54 deletions.
diff --git a/DragaliaAPI/DragaliaAPI.Integration.Test/Features/Trade/TreasureTradeTest.cs b/DragaliaAPI/DragaliaAPI.Integration.Test/Features/Trade/TreasureTradeTest.cs
@@ -113,4 +113,38 @@ await Client.PostMsgpack<TreasureTradeTradeResponse>(
         response.UpdateDataList.Should().NotBeNull();
         response.UpdateDataList.WeaponSkinList.Should().Contain(x => x.WeaponSkinId == 30159921);
     }
+
+    [Fact]
+    public async Task Trade_MultiDragon_AlmostFull_Trades()
+    {
+        this.ApiContext.PlayerDragonData.ExecuteDelete();
+        this.ApiContext.PlayerUserData.ExecuteUpdate(e =>
+            e.SetProperty(x => x.MaxDragonQuantity, 2)
+        );
+
+        int highBrunhildaTrade = 10030302;
+
+        TreasureTradeTradeResponse response = (
+            await Client.PostMsgpack<TreasureTradeTradeResponse>(
+                "treasure_trade/trade",
+                new TreasureTradeTradeRequest(1003, highBrunhildaTrade, null, 4)
+            )
+        ).Data;
+
+        response.UpdateDataList.DragonList.Should().HaveCount(2);
+        response.EntityResult.OverPresentEntityList.Should().HaveCount(2);
+
+        this.ApiContext.PlayerDragonData.ToList()
+            .Should()
+            .HaveCount(2)
+            .And.AllSatisfy(x => x.DragonId.Should().Be(Dragons.HighBrunhilda));
+        this.ApiContext.PlayerPresents.ToList()
+            .Should()
+            .HaveCount(2)
+            .And.AllSatisfy(x =>
+            {
+                x.EntityType.Should().Be(EntityTypes.Dragon);
+                x.EntityId.Should().Be((int)Dragons.HighBrunhilda);
+            });
+    }
 }
diff --git a/DragaliaAPI/DragaliaAPI/Features/Present/IPresentService.cs b/DragaliaAPI/DragaliaAPI/Features/Present/IPresentService.cs
@@ -12,4 +12,5 @@ public interface IPresentService
     void AddPresent(Present present);
 
     void AddPresent(IEnumerable<Present> presents);
+    IEnumerable<AtgenBuildEventRewardEntityList> GetTrackedPresentList();
 }
diff --git a/DragaliaAPI/DragaliaAPI/Features/Present/PresentService.cs b/DragaliaAPI/DragaliaAPI/Features/Present/PresentService.cs
@@ -1,3 +1,4 @@
+using DragaliaAPI.Database;
 using DragaliaAPI.Models.Generated;
 using DragaliaAPI.Shared.PlayerDetails;
 using Microsoft.EntityFrameworkCore;
@@ -7,28 +8,20 @@ namespace DragaliaAPI.Features.Present;
 /// <summary>
 /// Base present service to be used by other features to check/add present data.
 /// </summary>
-public class PresentService : IPresentService
+public class PresentService(
+    IPresentRepository presentRepository,
+    IPlayerIdentityService playerIdentityService,
+    ApiContext apiContext
+) : IPresentService
 {
-    private readonly IPresentRepository presentRepository;
-    private readonly IPlayerIdentityService playerIdentityService;
-
-    public PresentService(
-        IPresentRepository presentRepository,
-        IPlayerIdentityService playerIdentityService
-    )
-    {
-        this.presentRepository = presentRepository;
-        this.playerIdentityService = playerIdentityService;
-    }
-
     public async Task<PresentNotice> GetPresentNotice()
     {
         return new()
         {
-            PresentCount = await this.presentRepository.Presents.CountAsync(x =>
+            PresentCount = await presentRepository.Presents.CountAsync(x =>
                 x.ReceiveLimitTime == null
             ),
-            PresentLimitCount = await this.presentRepository.Presents.CountAsync(x =>
+            PresentLimitCount = await presentRepository.Presents.CountAsync(x =>
                 x.ReceiveLimitTime != null
             ),
         };
@@ -41,8 +34,15 @@ public void AddPresent(Present present)
 
     public void AddPresent(IEnumerable<Present> presents)
     {
-        this.presentRepository.AddPlayerPresents(
-            presents.Select(x => x.ToEntity(this.playerIdentityService.ViewerId))
+        presentRepository.AddPlayerPresents(
+            presents.Select(x => x.ToEntity(playerIdentityService.ViewerId))
         );
     }
+
+    public IEnumerable<AtgenBuildEventRewardEntityList> GetTrackedPresentList() =>
+        apiContext.PlayerPresents.Local.Select(x => new AtgenBuildEventRewardEntityList(
+            x.EntityType,
+            x.EntityId,
+            x.EntityQuantity
+        ));
 }
diff --git a/DragaliaAPI/DragaliaAPI/Features/Reward/Handlers/DragonHandler.cs b/DragaliaAPI/DragaliaAPI/Features/Reward/Handlers/DragonHandler.cs
@@ -16,6 +16,14 @@ ILogger<DragonHandler> logger
 
     public async Task<GrantReturn> Grant(Entity entity)
     {
+        if (entity.Quantity > 1)
+        {
+            throw new ArgumentException(
+                "Cannot process dragons with quantity >1 in single reward handler due to the possibility of multiple results. Use the batch handler instead.",
+                nameof(entity)
+            );
+        }
+
         Dragons dragon = (Dragons)entity.Id;
         if (!Enum.IsDefined(dragon))
         {

diff --git a/DragaliaAPI/DragaliaAPI/Features/Reward/IRewardService.cs b/DragaliaAPI/DragaliaAPI/Features/Reward/IRewardService.cs
@@ -37,7 +37,7 @@ int atk
     EntityResult GetEntityResult();
     IEnumerable<ConvertedEntity> GetConvertedEntityList();
 
-    Task<IDictionary<TKey, RewardGrantResult>> GrantRewardsWithBatchHandler<TKey>(
+    Task<IDictionary<TKey, RewardGrantResult>> BatchGrantRewards<TKey>(
         IDictionary<TKey, Entity> entities
     )
         where TKey : struct;

diff --git a/DragaliaAPI/DragaliaAPI/Features/Reward/RewardService.cs b/DragaliaAPI/DragaliaAPI/Features/Reward/RewardService.cs
@@ -48,7 +48,7 @@ public async Task GrantRewards(IEnumerable<Entity> entities)
         }
     }
 
-    public async Task<IDictionary<TKey, RewardGrantResult>> GrantRewardsWithBatchHandler<TKey>(
+    public async Task<IDictionary<TKey, RewardGrantResult>> BatchGrantRewards<TKey>(
         IDictionary<TKey, Entity> entities
     )
         where TKey : struct
@@ -63,23 +63,30 @@ IDictionary<TKey, Entity> entities
         foreach ((EntityTypes type, Dictionary<TKey, Entity> dictionary) in grouping)
         {
             if (
-                batchRewardHandlers.FirstOrDefault(x => x.SupportedTypes.Contains(type))
-                is not { } batchRewardHandler
+                batchRewardHandlers.FirstOrDefault(x => x.SupportedTypes.Contains(type)) is
+                { } batchRewardHandler
             )
             {
-                throw new NotSupportedException(
-                    $"Entities of type {type} are not supported in this method"
+                IDictionary<TKey, GrantReturn> batchResult = await batchRewardHandler.GrantRange(
+                    dictionary
                 );
-            }
-
-            IDictionary<TKey, GrantReturn> batchResult = await batchRewardHandler.GrantRange(
-                dictionary
-            );
 
-            foreach ((TKey key, GrantReturn grantReturn) in batchResult)
+                foreach ((TKey key, GrantReturn grantReturn) in batchResult)
+                {
+                    await ProcessGrantResult(grantReturn, dictionary[key]);
+                    result.Add(key, grantReturn.Result);
+                }
+            }
+            else
             {
-                await ProcessGrantResult(grantReturn, dictionary[key]);
-                result.Add(key, grantReturn.Result);
+                IRewardHandler handler = this.GetHandler(type);
+
+                foreach ((TKey key, Entity entity) in dictionary)
+                {
+                    GrantReturn grantReturn = await handler.Grant(entity);
+                    await ProcessGrantResult(grantReturn, entity);
+                    result.Add(key, grantReturn.Result);
+                }
             }
         }
 

diff --git a/DragaliaAPI/DragaliaAPI/Features/Summoning/SummonService.cs b/DragaliaAPI/DragaliaAPI/Features/Summoning/SummonService.cs
@@ -479,13 +479,9 @@ EntityResult EntityResult
             returnedResult.Add(processedResult);
         }
 
-        IEnumerable<AtgenBuildEventRewardEntityList> overPresentEntityList = apiContext
-            .PlayerPresents.Local.Where(x => x.EntityType == EntityTypes.Dragon)
-            .Select(x => new AtgenBuildEventRewardEntityList(
-                x.EntityType,
-                x.EntityId,
-                x.EntityQuantity
-            ));
+        IEnumerable<AtgenBuildEventRewardEntityList> overPresentEntityList = presentService
+            .GetTrackedPresentList()
+            .Where(x => x.EntityType == EntityTypes.Dragon);
 
         return (
             returnedResult,

diff --git a/DragaliaAPI/DragaliaAPI/Features/Summoning/UnitService.cs b/DragaliaAPI/DragaliaAPI/Features/Summoning/UnitService.cs
@@ -34,7 +34,7 @@ public async Task<IList<CharaNewCheckResult>> AddCharas(IList<Charas> idList)
             .ToDictionary();
 
         IDictionary<int, RewardGrantResult> outputRewardDict =
-            await rewardService.GrantRewardsWithBatchHandler(inputRewardDict);
+            await rewardService.BatchGrantRewards(inputRewardDict);
 
         List<CharaNewCheckResult> result = [];
 
@@ -55,7 +55,7 @@ public async Task<IList<DragonNewCheckResult>> AddDragons(List<Dragons> idList)
             .ToDictionary();
 
         IDictionary<int, RewardGrantResult> outputRewardDict =
-            await rewardService.GrantRewardsWithBatchHandler(inputRewardDict);
+            await rewardService.BatchGrantRewards(inputRewardDict);
 
         IEnumerable<Present.Present> presentsToAdd = outputRewardDict
             .Where(kvp => kvp.Value == RewardGrantResult.Limit)

diff --git a/DragaliaAPI/DragaliaAPI/Features/Trade/ITradeService.cs b/DragaliaAPI/DragaliaAPI/Features/Trade/ITradeService.cs
@@ -20,4 +20,5 @@ Task DoTrade(
         IEnumerable<AtgenNeedUnitList>? needUnitList = null
     );
     Task DoAbilityCrestTrade(int id, int count);
+    EntityResult GetEntityResult();
 }
diff --git a/DragaliaAPI/DragaliaAPI/Features/Trade/TradeService.cs b/DragaliaAPI/DragaliaAPI/Features/Trade/TradeService.cs
@@ -164,24 +164,53 @@ await paymentService.ProcessPayment(
             );
         }
 
-        RewardGrantResult result = await rewardService.GrantReward(
-            new(
-                trade.DestinationEntityType,
-                trade.DestinationEntityId,
-                trade.DestinationEntityQuantity * count,
-                trade.DestinationLimitBreakCount
-            )
-        );
+        Dictionary<int, Entity> entities;
 
-        if (result == RewardGrantResult.Limit)
+        if (trade.DestinationEntityType == EntityTypes.Dragon)
+        {
+            // We must flatten dragons out into multiple entities, since each grant could have a different result
+            entities = Enumerable
+                .Repeat(
+                    new Entity(
+                        trade.DestinationEntityType,
+                        trade.DestinationEntityId,
+                        1,
+                        trade.DestinationLimitBreakCount
+                    ),
+                    trade.DestinationEntityQuantity * count
+                )
+                .Select((x, index) => KeyValuePair.Create(index, x))
+                .ToDictionary();
+        }
+        else
         {
-            presentService.AddPresent(
-                new Present.Present(
-                    PresentMessage.TreasureTrade,
+            entities = new()
+            {
+                [1] = new Entity(
                     trade.DestinationEntityType,
-                    trade.DestinationEntityId
+                    trade.DestinationEntityId,
+                    trade.DestinationEntityQuantity * count,
+                    trade.DestinationLimitBreakCount
                 )
-            );
+            };
+        }
+
+        IDictionary<int, RewardGrantResult> batchResult = await rewardService.BatchGrantRewards(
+            entities
+        );
+
+        foreach ((_, RewardGrantResult result) in batchResult)
+        {
+            if (result == RewardGrantResult.Limit)
+            {
+                presentService.AddPresent(
+                    new Present.Present(
+                        PresentMessage.TreasureTrade,
+                        trade.DestinationEntityType,
+                        trade.DestinationEntityId
+                    )
+                );
+            }
         }
 
         await tradeRepository.AddTrade(tradeType, tradeId, count, DateTimeOffset.UtcNow);
@@ -217,4 +246,12 @@ await rewardService.GrantReward(
 
         await tradeRepository.AddTrade(TradeType.AbilityCrest, id, count);
     }
+
+    public EntityResult GetEntityResult()
+    {
+        EntityResult result = rewardService.GetEntityResult();
+        result.OverPresentEntityList = presentService.GetTrackedPresentList();
+
+        return result;
+    }
 }
diff --git a/DragaliaAPI/DragaliaAPI/Features/Trade/TreasureTradeController.cs b/DragaliaAPI/DragaliaAPI/Features/Trade/TreasureTradeController.cs
@@ -44,6 +44,7 @@ await tradeService.DoTrade(
         resp.UpdateDataList = await updateDataService.SaveChangesAsync(cancellationToken);
         resp.TreasureTradeAllList = tradeService.GetCurrentTreasureTradeList();
         resp.UserTreasureTradeList = await tradeService.GetUserTreasureTradeList();
+        resp.EntityResult = tradeService.GetEntityResult();
 
         return Ok(resp);
     }