updates

Author: taylordowns2000

packages/lightning-mock/src/api-dev.ts

@@ -201,7 +201,12 @@ const setupDevAPI = (
       }) => {
         if (evt.runId === runId) {
           state.events.removeListener(RUN_COMPLETE, handler);
-          resolve(evt.payload.final_state);
+          const { final_dataclip_id, final_state } = evt.payload;
+          if (final_dataclip_id) {
+            resolve(state.dataclips?.[final_dataclip_id] ?? null);
+          } else {
+            resolve(final_state);
+          }
         }
       };
       state.events.addListener(RUN_COMPLETE, handler);

packages/lightning-mock/src/api-sockets.ts

@@ -393,6 +393,7 @@ const createSocketAPI = (
   ) {
     const { ref, join_ref, topic } = evt;
     const {
+      final_dataclip_id,
       final_state,
       reason,
       error_type,
@@ -410,7 +411,9 @@ const createSocketAPI = (
     if (!state.results[runId]) {
       state.results[runId] = { state: null, workerId: 'mock' };
     }
-    if (final_state) {
+    if (final_dataclip_id) {
+      state.results[runId].state = state.dataclips?.[final_dataclip_id] ?? null;
+    } else if (final_state) {
       state.results[runId].state = final_state;
     }
 

packages/ws-worker/test/lightning.test.ts

@@ -291,8 +291,11 @@ test.serial(
     return new Promise((done) => {
       const run = getRun();
       lng.onSocketEvent(e.RUN_COMPLETE, run.id, (evt: any) => {
-        const { final_state } = evt.payload;
-        t.assert(typeof final_state === 'object');
+        const { final_dataclip_id, final_state } = evt.payload;
+        t.assert(
+          typeof final_dataclip_id === 'string' ||
+            typeof final_state === 'object'
+        );
         t.pass('run complete event received');
         done();
       });