fix: address PR comments

OasisDEX · May 27, 2024 · e374960 · e374960
1 parent e2e9b4b
commit e374960
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 12 deletions.
diff --git a/sdk/order-planner-common/src/implementation/OrderPlanner.ts b/sdk/order-planner-common/src/implementation/OrderPlanner.ts
@@ -1,9 +1,15 @@
 import { Order, type IPositionsManager, TransactionInfo } from '@summerfi/sdk-common/orders'
-import { ISimulation, SimulationType, steps } from '@summerfi/sdk-common/simulation'
+import {
+  ISimulation,
+  SimulationSteps,
+  SimulationType,
+  steps,
+} from '@summerfi/sdk-common/simulation'
 import { Maybe } from '@summerfi/sdk-common/common'
 import {
   ActionBuildersMap,
   ActionCall,
+  FilterStep,
   IActionBuilder,
   IStepBuilderContext,
   StepBuilderContext,
@@ -34,7 +40,7 @@ export class OrderPlanner implements IOrderPlanner {
     context.startSubContext()
 
     for (const step of simulation.steps) {
-      const stepBuilder = this._getActionBuilder(actionBuildersMap, step)
+      const stepBuilder = this._getActionBuilder(actionBuildersMap, step.type)
       if (!stepBuilder) {
         throw new Error(`No step builder found for step type ${step.type}`)
       }
@@ -68,17 +74,17 @@ export class OrderPlanner implements IOrderPlanner {
     })
   }
 
-  private _getActionBuilder<StepType extends steps.Steps>(
-    actionBuildersMap: ActionBuildersMap,
-    step: StepType,
-  ): Maybe<IActionBuilder<StepType>> {
-    const builder = actionBuildersMap[step.type]
+  private _getActionBuilder<
+    StepType extends SimulationSteps,
+    Step extends FilterStep<StepType, steps.Steps>,
+  >(actionBuildersMap: ActionBuildersMap, stepType: StepType): Maybe<IActionBuilder<Step>> {
+    const builder = actionBuildersMap[stepType]
 
     if (!builder) {
       return undefined
     }
 
-    return new builder() as IActionBuilder<StepType>
+    return new builder()
   }
 
   private async _generateOrder(params: {

diff --git a/sdk/protocol-plugins/src/implementation/BaseProtocolPlugin.ts b/sdk/protocol-plugins/src/implementation/BaseProtocolPlugin.ts
@@ -158,7 +158,7 @@ export abstract class BaseProtocolPlugin implements IProtocolPlugin {
       return undefined
     }
 
-    return new BuilderClass() as IActionBuilder<Step>
+    return new BuilderClass()
   }
 
   /** HELPERS */

diff --git a/sdk/protocol-plugins/tests/utils/ProtocolPluginMock.ts b/sdk/protocol-plugins/tests/utils/ProtocolPluginMock.ts
@@ -117,7 +117,7 @@ export class ProtocolPluginMock implements IProtocolPlugin {
       return undefined
     }
 
-    return new builder() as IActionBuilder<Step>
+    return new builder()
   }
 
   async getImportPositionTransaction(params: {
@@ -167,7 +167,7 @@ export class EmptyProtocolPluginMock implements IProtocolPlugin {
       return undefined
     }
 
-    return new builder() as IActionBuilder<Step>
+    return new builder()
   }
 
   async getImportPositionTransaction(params: {
@@ -222,7 +222,7 @@ export class NoCheckpointProtocolPluginMock implements IProtocolPlugin {
       return undefined
     }
 
-    return new builder() as IActionBuilder<Step>
+    return new builder()
   }
 
   async getImportPositionTransaction(params: {