From a08634a1e832f6aae83d72a365006d3649f2c979 Mon Sep 17 00:00:00 2001
From: Jethary <jethary@opentrons.com>
Date: Mon, 10 Feb 2025 15:10:24 -0500
Subject: [PATCH] fix unit test

---
 .../file-data/__tests__/createFile.test.ts    | 20 +++++++++++++------
 .../src/file-data/selectors/pythonFile.ts     |  8 ++++----
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/protocol-designer/src/file-data/__tests__/createFile.test.ts b/protocol-designer/src/file-data/__tests__/createFile.test.ts
index 527e1cccf20..4b3ad1103ca 100644
--- a/protocol-designer/src/file-data/__tests__/createFile.test.ts
+++ b/protocol-designer/src/file-data/__tests__/createFile.test.ts
@@ -70,6 +70,12 @@ describe('createFile selector', () => {
   afterEach(() => {
     vi.restoreAllMocks()
   })
+  const entities = {
+    moduleEntities: v7Fixture.moduleEntities,
+    labwareEntities,
+    pipetteEntities,
+    liquidEntities: ingredients,
+  }
   it('should return a schema-valid JSON V8 protocol', () => {
     // @ts-expect-error(sa, 2021-6-15): resultFunc not part of Selector type
     const result = createFile.resultFunc(
@@ -78,16 +84,13 @@ describe('createFile selector', () => {
       v7Fixture.robotStateTimeline,
       OT2_ROBOT_TYPE,
       dismissedWarnings,
-      ingredients,
       ingredLocations,
       v7Fixture.savedStepForms,
       v7Fixture.orderedStepIds,
-      labwareEntities,
-      v7Fixture.moduleEntities,
-      pipetteEntities,
       labwareNicknamesById,
       labwareDefsByURI,
-      {}
+      {},
+      entities
     )
     expectResultToMatchSchema(result)
 
@@ -99,7 +102,12 @@ describe('createFile selector', () => {
 
   it('should return a valid Python protocol file', () => {
     // @ts-expect-error(sa, 2021-6-15): resultFunc not part of Selector type
-    const result = createPythonFile.resultFunc(fileMetadata, OT2_ROBOT_TYPE, {})
+    const result = createPythonFile.resultFunc(
+      fileMetadata,
+      OT2_ROBOT_TYPE,
+      entities,
+      v7Fixture.initialRobotState
+    )
     // This is just a quick smoke test to make sure createPythonFile() produces
     // something that looks like a Python file. The individual sections of the
     // generated Python will be tested in separate unit tests.
diff --git a/protocol-designer/src/file-data/selectors/pythonFile.ts b/protocol-designer/src/file-data/selectors/pythonFile.ts
index eb6daa6d116..1fc6c3cfe47 100644
--- a/protocol-designer/src/file-data/selectors/pythonFile.ts
+++ b/protocol-designer/src/file-data/selectors/pythonFile.ts
@@ -76,10 +76,10 @@ export function pythonDefRun(
   invariantContext: InvariantContext,
   robotState: TimelineFrame
 ): string {
-  const loadModules = getLoadModules(
-    invariantContext.moduleEntities,
-    robotState.modules
-  )
+  const { moduleEntities } = invariantContext
+
+  const loadModules = getLoadModules(moduleEntities, robotState.modules)
+
   const sections: string[] = [
     ...loadModules,
     // loadLabware(),