feat(step-generation): check for shuttle errors and mismatched labware

rclarke0 · rclarke0 · commit f2f0b124704f · 2025-12-08T17:27:55.000-05:00
diff --git a/step-generation/src/__tests__/flexStackerRetrieve.test.ts b/step-generation/src/__tests__/flexStackerRetrieve.test.ts
@@ -13,9 +13,14 @@ import {
   getInitialRobotStateStandard,
   makeContext,
 } from '../fixtures'
+import { flexStackerStateGetter } from '../robotStateSelectors'
 
 import type { LabwareDefinition2 } from '@opentrons/shared-data'
-import type { InvariantContext, RobotState } from '../types'
+import type {
+  FlexStackerModuleState,
+  InvariantContext,
+  RobotState,
+} from '../types'
 
 const mockLabwareId = 'labwareId'
 const mockLabwareId2 = 'labwareId2'
@@ -35,6 +40,9 @@ describe('flexStackerRetrieve', () => {
       model: FLEX_STACKER_MODULE_V1,
       pythonName: 'mock_flex_stacker_1',
     }
+    vi.mocked(flexStackerStateGetter).mockReturnValue(
+      {} as FlexStackerModuleState
+    )
   })
   it('creates flex stacker retrieve command', () => {
     robotState = {
@@ -117,7 +125,7 @@ describe('flexStackerRetrieve', () => {
       },
       labware: {
         [mockLabwareId]: {
-          stack: ['D2'],
+          stack: [mockLabwareId, 'D2'],
         },
       },
     }
@@ -143,4 +151,52 @@ describe('flexStackerRetrieve', () => {
       type: 'HOPPER_EMPTY',
     })
   })
+  it('raises an error if the shuttle is full', () => {
+    vi.mocked(flexStackerStateGetter).mockReturnValue({
+      labwareOnShuttle: {
+        primaryLabwareId: mockLabwareId,
+        adapterLabwareId: null,
+        lidLabwareId: null,
+      },
+      labwareInHopper: null,
+      maxPoolCount: 10,
+      storedLabwareDetails: null,
+      type: FLEX_STACKER_MODULE_TYPE,
+    })
+    robotState = {
+      ...robotState,
+      modules: {
+        [mockModuleId]: {
+          slot: 'D3',
+          moduleState: {} as any,
+        },
+      },
+      labware: {
+        [mockLabwareId]: {
+          stack: [mockLabwareId, 'D3'],
+        },
+      },
+    }
+    invariantContext = {
+      ...invariantContext,
+      labwareEntities: {
+        [mockLabwareId]: {
+          labwareDefURI: 'mockURI',
+          def: fixture96Plate as LabwareDefinition2,
+          pythonName: 'wellPlate_1',
+          id: mockLabwareId,
+        },
+      },
+    }
+    const result = flexStackerRetrieve(
+      {
+        moduleId: mockModuleId,
+      },
+      invariantContext,
+      robotState
+    )
+    expect(getErrorResult(result).errors[0]).toMatchObject({
+      type: 'SHUTTLE_FULL',
+    })
+  })
 })
diff --git a/step-generation/src/__tests__/flexStackerStore.test.ts b/step-generation/src/__tests__/flexStackerStore.test.ts
@@ -1,14 +1,26 @@
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 import {
+  fixture96Plate,
+  fixture384Plate,
   FLEX_STACKER_MODULE_TYPE,
   FLEX_STACKER_MODULE_V1,
 } from '@opentrons/shared-data'
 
 import { flexStackerStore } from '../commandCreators/atomic/flexStackerStore'
-import { getInitialRobotStateStandard, makeContext } from '../fixtures'
+import {
+  getErrorResult,
+  getInitialRobotStateStandard,
+  makeContext,
+} from '../fixtures'
+import { flexStackerStateGetter } from '../robotStateSelectors'
 
-import type { InvariantContext, RobotState } from '../types'
+import type { LabwareDefinition2 } from '@opentrons/shared-data'
+import type {
+  FlexStackerModuleState,
+  InvariantContext,
+  RobotState,
+} from '../types'
 
 const moduleId = 'flexStackerId'
 vi.mock('../robotStateSelectors')
@@ -25,13 +37,60 @@ describe('flexStackerStore', () => {
       model: FLEX_STACKER_MODULE_V1,
       pythonName: 'mock_flex_stacker_1',
     }
+    invariantContext.labwareEntities = {
+      mockLabwareId: {
+        id: 'mockLabwareId',
+        labwareDefURI: 'mockURI',
+        def: fixture96Plate as LabwareDefinition2,
+        pythonName: 'wellPlate_1',
+      },
+    }
+    vi.mocked(flexStackerStateGetter).mockReturnValue({
+      labwareOnShuttle: {
+        primaryLabwareId: 'mockLabwareId',
+        adapterLabwareId: null,
+        lidLabwareId: null,
+      },
+      labwareInHopper: [
+        {
+          primaryLabwareId: 'mockLabwareId',
+          adapterLabwareId: null,
+          lidLabwareId: null,
+        },
+      ],
+      maxPoolCount: 10,
+      storedLabwareDetails: {
+        moduleId,
+        initialCount: 5,
+        primaryLabware: {
+          loadName: 'fixture_96_plate',
+          namespace: 'opentrons',
+          version: 1,
+        },
+        lidLabware: null,
+        adapterLabware: null,
+      },
+      type: FLEX_STACKER_MODULE_TYPE,
+    } as FlexStackerModuleState)
+    robotState = {
+      ...robotState,
+      modules: {
+        [moduleId]: {
+          slot: 'D3',
+          moduleState: {} as FlexStackerModuleState,
+        },
+      },
+      labware: {
+        mockLabwareId: {
+          stack: ['mockLabwareId', 'D3'],
+        },
+      },
+    }
   })
+
   it('creates flex stacker store command', () => {
     const result = flexStackerStore(
-      {
-        moduleId,
-        strategy: 'automatic',
-      },
+      { moduleId, strategy: 'automatic' },
       invariantContext,
       robotState
     )
@@ -49,4 +108,114 @@ describe('flexStackerStore', () => {
       python: 'mock_flex_stacker_1.store()',
     })
   })
+
+  it('raises an error if the shuttle is empty', () => {
+    vi.mocked(flexStackerStateGetter).mockReturnValue({
+      labwareOnShuttle: null,
+      labwareInHopper: null,
+      maxPoolCount: 10,
+      storedLabwareDetails: null,
+      type: FLEX_STACKER_MODULE_TYPE,
+    })
+
+    const emptyShuttleRobotState: RobotState = {
+      ...robotState,
+      modules: {
+        [moduleId]: {
+          slot: 'D3',
+          moduleState: {} as FlexStackerModuleState,
+        },
+      },
+    }
+
+    const emptyShuttleInvariantContext: InvariantContext = {
+      ...invariantContext,
+      labwareEntities: {
+        [moduleId]: {
+          id: moduleId,
+          labwareDefURI: 'mockURI',
+          def: fixture96Plate as LabwareDefinition2,
+          pythonName: 'wellPlate_1',
+        },
+      },
+    }
+
+    const result = flexStackerStore(
+      { moduleId, strategy: 'automatic' },
+      emptyShuttleInvariantContext,
+      emptyShuttleRobotState
+    )
+
+    expect(getErrorResult(result).errors[0]).toMatchObject({
+      type: 'SHUTTLE_EMPTY',
+    })
+  })
+
+  it('raises an error if the labware to be stored does not match the current labware in the hopper', () => {
+    vi.mocked(flexStackerStateGetter).mockReturnValue({
+      labwareOnShuttle: {
+        primaryLabwareId: 'labware-in-shuttle',
+        adapterLabwareId: null,
+        lidLabwareId: null,
+      },
+      labwareInHopper: [
+        {
+          primaryLabwareId: 'labware-in-hopper',
+          adapterLabwareId: null,
+          lidLabwareId: null,
+        },
+      ],
+      maxPoolCount: 10,
+      storedLabwareDetails: {
+        moduleId,
+        initialCount: 5,
+        primaryLabware: {
+          loadName: 'wellPlate_1',
+          namespace: 'opentrons',
+          version: 1,
+        },
+        lidLabware: null,
+        adapterLabware: null,
+      },
+      type: FLEX_STACKER_MODULE_TYPE,
+    })
+    const mismatchRobotState: RobotState = {
+      ...robotState,
+      modules: {
+        [moduleId]: {
+          slot: 'D3',
+          moduleState: {} as FlexStackerModuleState,
+        },
+      },
+      labware: {
+        'labware-in-hopper': { stack: ['labware-in-hopper', 'D3'] },
+        'labware-in-shuttle': { stack: ['labware-in-shuttle', 'D3'] },
+      },
+    }
+    const mismatchInvariantContext: InvariantContext = {
+      ...invariantContext,
+      labwareEntities: {
+        'labware-in-hopper': {
+          id: 'labware-in-hopper',
+          labwareDefURI: 'mockURI',
+          def: fixture96Plate as LabwareDefinition2,
+          pythonName: 'wellPlate_1',
+        },
+        'labware-in-shuttle': {
+          id: 'labware-in-shuttle',
+          labwareDefURI: 'differentMockURI',
+          def: fixture384Plate as LabwareDefinition2,
+          pythonName: 'wellPlate_2',
+        },
+      },
+    }
+    const result = flexStackerStore(
+      { moduleId, strategy: 'automatic' },
+      mismatchInvariantContext,
+      mismatchRobotState
+    )
+    expect(getErrorResult(result).errors[0]).toMatchObject({
+      type: 'MISMATCHED_STACKER_LABWARE_TYPE',
+    })
+  })
 })
diff --git a/step-generation/src/commandCreators/atomic/flexStackerRetrieve.ts b/step-generation/src/commandCreators/atomic/flexStackerRetrieve.ts
@@ -1,5 +1,6 @@
 import * as errorCreators from '../../errorCreators'
-import { getLabwareIdOnHopper, uuid } from '../../utils'
+import { flexStackerStateGetter } from '../../robotStateSelectors'
+import { getLabwareIdOnHopper, getLabwareIdOnShuttle, uuid } from '../../utils'
 
 import type { FlexStackerRetrieveCreateCommand } from '@opentrons/shared-data'
 import type { CommandCreator } from '../../types'
@@ -12,15 +13,20 @@ export const flexStackerRetrieve: CommandCreator<
   const { moduleEntities, labwareEntities } = invariantContext
   const modulePythonName = moduleEntities[moduleId].pythonName
   const moduleLocation = modules[moduleId].slot
-  const labwareIdOnModule = getLabwareIdOnHopper(labware, moduleLocation)
-  const labwarePythonName = labwareEntities[labwareIdOnModule]?.pythonName
-  // TODO: add error for if there is labware on the shuttle
-  if (!labwareIdOnModule) {
+  const flexStackerState = flexStackerStateGetter(robotState, moduleId)
+  const labwareIdOnHopper = getLabwareIdOnHopper(labware, moduleLocation)
+  const labwarePythonName = labwareEntities[labwareIdOnHopper]?.pythonName
+
+  if (flexStackerState !== null && getLabwareIdOnShuttle(flexStackerState)) {
+    return {
+      errors: [errorCreators.flexStackerShuttleFull()],
+    }
+  }
+  if (!labwareIdOnHopper) {
     return {
       errors: [errorCreators.flexStackerHopperEmpty()],
     }
   }
-
   return {
     commands: [
       {
diff --git a/step-generation/src/commandCreators/atomic/flexStackerStore.ts b/step-generation/src/commandCreators/atomic/flexStackerStore.ts
@@ -1,13 +1,39 @@
-import { uuid } from '../../utils'
+import * as errorCreators from '../../errorCreators'
+import { flexStackerStateGetter } from '../../robotStateSelectors'
+import {
+  getLabwareIdOnShuttle,
+  labwareMatchesLabwareInHopper,
+  uuid,
+} from '../../utils'
 
 import type { FlexStackerStoreCreateCommand } from '@opentrons/shared-data'
 import type { CommandCreator } from '../../types'
 
 export const flexStackerStore: CommandCreator<
   FlexStackerStoreCreateCommand['params']
-> = (args, invariantContext) => {
+> = (args, invariantContext, robotState) => {
   const { moduleId } = args
   const pythonName = invariantContext.moduleEntities[moduleId].pythonName
+  const flexStackerState = flexStackerStateGetter(robotState, moduleId)
+  const labwareId = flexStackerState?.labwareOnShuttle?.primaryLabwareId ?? null
+  if (flexStackerState !== null && !getLabwareIdOnShuttle(flexStackerState)) {
+    return {
+      errors: [errorCreators.flexStackerShuttleEmpty()],
+    }
+  }
+
+  if (
+    labwareId !== null &&
+    !labwareMatchesLabwareInHopper(
+      labwareId,
+      invariantContext,
+      flexStackerState
+    )
+  ) {
+    return {
+      errors: [errorCreators.flexStackerLabwareTypeMismatch()],
+    }
+  }
   return {
     commands: [
       {
diff --git a/step-generation/src/utils/misc.ts b/step-generation/src/utils/misc.ts
