From a53dfe6200166564c7350d91b054f1029b0c1e02 Mon Sep 17 00:00:00 2001
From: Saga Dasgupta <saga.dasgupta@shopify.com>
Date: Thu, 19 Sep 2024 19:53:52 +0000
Subject: [PATCH] draftable dev works now and addressing pr comments

---
 .../models/extensions/extension-instance.ts   |  7 +--
 .../app/src/cli/services/build/extension.ts   |  6 +++
 .../cli/services/dev/update-extension.test.ts | 45 +++++++++++++++++++
 .../src/cli/services/dev/update-extension.ts  | 17 +++++--
 4 files changed, 66 insertions(+), 9 deletions(-)

diff --git a/packages/app/src/cli/models/extensions/extension-instance.ts b/packages/app/src/cli/models/extensions/extension-instance.ts
index addb7eba72..36048ab391 100644
--- a/packages/app/src/cli/models/extensions/extension-instance.ts
+++ b/packages/app/src/cli/models/extensions/extension-instance.ts
@@ -28,7 +28,7 @@ import {constantize, slugify} from '@shopify/cli-kit/common/string'
 import {hashString, randomUUID} from '@shopify/cli-kit/node/crypto'
 import {partnersFqdn} from '@shopify/cli-kit/node/context/fqdn'
 import {joinPath} from '@shopify/cli-kit/node/path'
-import {fileExists, touchFile, writeFile, copyFile} from '@shopify/cli-kit/node/fs'
+import {fileExists, touchFile, writeFile} from '@shopify/cli-kit/node/fs'
 import {getPathValue} from '@shopify/cli-kit/common/object'
 import {useThemebundling} from '@shopify/cli-kit/node/context/local'
 
@@ -336,13 +336,10 @@ export class ExtensionInstance<TConfiguration extends BaseConfigType = BaseConfi
 
     if (this.features.includes('bundling')) {
       // Modules that are going to be inclued in the bundle should be built in the bundle directory
-      this.outputPath = this.isFunctionExtension ? this.outputPath : joinPath(bundleDirectory, extensionId, outputFile)
+      this.outputPath = joinPath(bundleDirectory, extensionId, outputFile)
     }
 
     await this.build(options)
-    if (this.isFunctionExtension) {
-      await copyFile(this.outputPath, joinPath(bundleDirectory, extensionId, 'index.wasm'))
-    }
     if (this.isThemeExtension && useThemebundling()) {
       await bundleThemeExtension(this, options)
     }
diff --git a/packages/app/src/cli/services/build/extension.ts b/packages/app/src/cli/services/build/extension.ts
index bf81e275fb..77391cdf18 100644
--- a/packages/app/src/cli/services/build/extension.ts
+++ b/packages/app/src/cli/services/build/extension.ts
@@ -10,6 +10,7 @@ import {AbortError, AbortSilentError} from '@shopify/cli-kit/node/error'
 import lockfile from 'proper-lockfile'
 import {joinPath} from '@shopify/cli-kit/node/path'
 import {outputDebug} from '@shopify/cli-kit/node/output'
+import {readFile, touchFile, writeFile} from '@shopify/cli-kit/node/fs'
 import {Writable} from 'stream'
 
 export interface ExtensionBuildOptions {
@@ -141,11 +142,16 @@ export async function buildFunctionExtension(
   }
 
   try {
+    const bundlePath = joinPath(extension.outputPath.split('/').slice(0, -2).join('/'), 'index.wasm')
+    extension.outputPath = joinPath(extension.directory, joinPath('dist', 'function.wasm'))
     if (extension.isJavaScript) {
       await runCommandOrBuildJSFunction(extension, options)
     } else {
       await buildOtherFunction(extension, options)
     }
+    const base64Contents = await readFile(extension.outputPath, {encoding: 'base64'})
+    await touchFile(bundlePath)
+    await writeFile(bundlePath, base64Contents)
   } finally {
     await releaseLock()
   }
diff --git a/packages/app/src/cli/services/dev/update-extension.test.ts b/packages/app/src/cli/services/dev/update-extension.test.ts
index dd9f7a4837..b062684f66 100644
--- a/packages/app/src/cli/services/dev/update-extension.test.ts
+++ b/packages/app/src/cli/services/dev/update-extension.test.ts
@@ -1,6 +1,7 @@
 import {reloadExtensionConfig, updateExtensionDraft} from './update-extension.js'
 import {
   placeholderAppConfiguration,
+  testFunctionExtension,
   testDeveloperPlatformClient,
   testPaymentExtensions,
   testThemeExtensions,
@@ -9,12 +10,14 @@ import {
 import {parseConfigurationFile, parseConfigurationObjectAgainstSpecification} from '../../models/app/loader.js'
 import {DeveloperPlatformClient} from '../../utilities/developer-platform-client.js'
 import {ExtensionUpdateDraftMutationVariables} from '../../api/graphql/partners/generated/update-draft.js'
+import {uploadWasmBlob} from '../deploy/upload.js'
 import {inTemporaryDirectory, mkdir, writeFile} from '@shopify/cli-kit/node/fs'
 import {outputInfo} from '@shopify/cli-kit/node/output'
 import {describe, expect, vi, test} from 'vitest'
 import {joinPath} from '@shopify/cli-kit/node/path'
 import {platformAndArch} from '@shopify/cli-kit/node/os'
 
+vi.mock('../deploy/upload.js')
 vi.mock('@shopify/cli-kit/node/output')
 vi.mock('../../models/app/loader.js', async () => {
   const actual: any = await vi.importActual('../../models/app/loader.js')
@@ -186,6 +189,48 @@ describe('updateExtensionDraft()', () => {
     })
   })
 
+  test('updates draft successfully for function app extension', async () => {
+    const developerPlatformClient: DeveloperPlatformClient = testDeveloperPlatformClient()
+    await inTemporaryDirectory(async (tmpDir) => {
+      const mockExtension = await testFunctionExtension({dir: tmpDir})
+
+      const filepath = 'index.wasm'
+      const content = 'test content'
+      const base64Content = Buffer.from(content).toString('base64')
+      await mkdir(joinPath(mockExtension.directory, 'dist'))
+      await writeFile(joinPath(mockExtension.directory, 'dist', filepath), content)
+      vi.mocked(uploadWasmBlob).mockResolvedValue({url: 'url', moduleId: 'moduleId'})
+
+      await updateExtensionDraft({
+        extension: mockExtension,
+        developerPlatformClient,
+        apiKey,
+        registrationId,
+        stdout,
+        stderr,
+        appConfiguration: placeholderAppConfiguration,
+      })
+
+      expect(developerPlatformClient.updateExtension).toHaveBeenCalledWith({
+        apiKey,
+        context: '',
+        handle: mockExtension.handle,
+        registrationId,
+        config: JSON.stringify({
+          title: 'test function extension',
+          module_id: 'moduleId',
+          description: 'description',
+          app_key: 'mock-api-key',
+          api_type: 'product_discounts',
+          api_version: '2022-07',
+          enable_creation_ui: true,
+          localization: {},
+          uploaded_files: {'index.wasm': base64Content},
+        }),
+      })
+    })
+  })
+
   test('handles user errors with stderr message', async () => {
     const errorResponse = {
       extensionUpdateDraft: {
diff --git a/packages/app/src/cli/services/dev/update-extension.ts b/packages/app/src/cli/services/dev/update-extension.ts
index 2a3308e024..4dbd12ebd0 100644
--- a/packages/app/src/cli/services/dev/update-extension.ts
+++ b/packages/app/src/cli/services/dev/update-extension.ts
@@ -50,12 +50,21 @@ export async function updateExtensionDraft({
     config = (await extension.deployConfig({apiKey, developerPlatformClient, appConfiguration})) || {}
   }
 
+  let compiledFiles: string | undefined
+  let draftableConfig: {[key: string]: unknown} = {
+    ...config,
+    serialized_script: encodedFile,
+  }
+  if (extension.isFunctionExtension) {
+    compiledFiles = await readFile(extension.outputPath, {encoding: 'base64'})
+    draftableConfig = {
+      ...draftableConfig,
+      uploaded_files: {'index.wasm': compiledFiles},
+    }
+  }
   const extensionInput: ExtensionUpdateDraftMutationVariables = {
     apiKey,
-    config: JSON.stringify({
-      ...config,
-      serialized_script: encodedFile,
-    }),
+    config: JSON.stringify(draftableConfig),
     handle: extension.handle,
     context: extension.contextValue,
     registrationId,