Merge pull request #5 from fourTheorem/cmg/fix-skip-behaviour

fix: skip behaviour and recursive copying

Author: chrismcg

src/bundling.ts

@@ -21,6 +21,7 @@ export const DEFAULT_ASSET_EXCLUDES = [
   'node_modules/',
   'cdk.out/',
   '.git/',
+  'cdk',
 ];
 
 interface BundlingCommandOptions {
@@ -71,7 +72,7 @@ export class Bundling {
     return Code.fromAsset(options.rootDir, {
       assetHashType: AssetHashType.SOURCE,
       exclude: HASHABLE_DEPENDENCIES_EXCLUDE,
-      bundling: options.skip ? undefined : new Bundling(options),
+      bundling: new Bundling(options),
     });
   }
 
@@ -92,13 +93,11 @@ export class Bundling {
       rootDir,
       workspacePackage,
       image,
-      runtime,
       commandHooks,
       assetExcludes = DEFAULT_ASSET_EXCLUDES,
-      architecture = Architecture.ARM_64,
     } = props;
 
-    const bundlingCommands = this.createBundlingCommands({
+    const bundlingCommands = props.skip ? [] : this.createBundlingCommands({
       rootDir,
       workspacePackage,
       assetExcludes,
@@ -107,15 +106,7 @@ export class Bundling {
       outputDir: AssetStaging.BUNDLING_OUTPUT_DIR,
     });
 
-    this.image =
-      image ??
-      DockerImage.fromBuild(path.resolve(__dirname, '..', 'resources'), {
-        buildArgs: {
-          ...props.buildArgs,
-          IMAGE: runtime.bundlingImage.image,
-        },
-        platform: architecture.dockerPlatform,
-      });
+    this.image = image ?? this.createDockerImage(props);
 
     this.command = props.command ?? [
       'bash',
@@ -133,6 +124,22 @@ export class Bundling {
     this.bundlingFileAccess = props.bundlingFileAccess;
   }
 
+  private createDockerImage(props: BundlingProps): DockerImage {
+    // If skip is true then don't call DockerImage.fromBuild as that calls dockerExec.
+    // Return a dummy object of the right type as it's not going to be used.
+    if (props.skip) {
+      return new DockerImage('skipped');
+    }
+
+    return DockerImage.fromBuild(path.resolve(__dirname, '..', 'resources'), {
+      buildArgs: {
+        ...props.buildArgs,
+        IMAGE: props.runtime.bundlingImage.image,
+      },
+      platform: (props.architecture ?? Architecture.ARM_64).dockerPlatform,
+    });
+  }
+
   private createBundlingCommands(options: BundlingCommandOptions): string[] {
     const excludeArgs = options.assetExcludes.map((exclude) => `--exclude="${exclude}"`);
     const workspacePackage = options.workspacePackage;

src/function.ts

@@ -75,10 +75,12 @@ export class PythonFunction extends Function {
       throw new Error('Only Python runtimes are supported');
     }
 
+    const skip = !Stack.of(scope).bundlingRequired;
+
     const code = Bundling.bundle({
       rootDir,
       runtime,
-      skip: !Stack.of(scope).bundlingRequired,
+      skip: skip,
       architecture,
       workspacePackage,
       ...props.bundling,
@@ -95,18 +97,23 @@ export class PythonFunction extends Function {
       handler: resolvedHandler,
     });
 
+    if (skip) {
+      return;
+    }
+
     const assetPath = ((this.node.defaultChild) as CfnFunction).getMetadata('aws:asset:path');
-    if (assetPath) { // TODO - remove - we always need one
-      const codePath = path.join(process.env.CDK_OUTDIR as string, assetPath);
+    if (!assetPath) {
+      return;
+    }
 
-      const pythonPaths = getPthFilePaths(codePath);
+    const codePath = path.join(process.env.CDK_OUTDIR as string, assetPath);
+    const pythonPaths = getPthFilePaths(codePath);
 
-      if (pythonPaths.length > 0) {
-        let pythonPathValue = environment.PYTHONPATH;
-        const addedPaths = pythonPaths.join(':');
-        pythonPathValue = pythonPathValue ? `${pythonPathValue}:${addedPaths}` : addedPaths;
-        this.addEnvironment('PYTHONPATH', pythonPathValue);
-      }
+    if (pythonPaths.length > 0) {
+      let pythonPathValue = environment.PYTHONPATH;
+      const addedPaths = pythonPaths.join(':');
+      pythonPathValue = pythonPathValue ? `${pythonPathValue}:${addedPaths}` : addedPaths;
+      this.addEnvironment('PYTHONPATH', pythonPathValue);
     }
   }
 }

test/function.test.ts

@@ -1,10 +1,12 @@
 import { exec } from 'node:child_process';
 import * as fs from 'node:fs/promises';
+import * as os from 'node:os';
 import * as path from 'node:path';
 import { promisify } from 'node:util';
 import { App, Stack } from 'aws-cdk-lib';
 import { Match, Template } from 'aws-cdk-lib/assertions';
 import { Architecture, Runtime } from 'aws-cdk-lib/aws-lambda';
+import * as cxapi from 'aws-cdk-lib/cx-api';
 import { PythonFunction } from '../src';
 const execAsync = promisify(exec);
 
@@ -27,9 +29,41 @@ async function getDockerHostArch(): Promise<Architecture> {
   }
 }
 
-test('Create a function from basic_app', async () => {
+/**
+ * Create a new CDK App and Stack with the given name and set the context to ensure
+ * that the 'aws:asset:path' metadata is set.
+ *
+ * @returns The App and Stack
+ */
+async function createStack(name = 'test'): Promise<{ app: App; stack: Stack }> {
   const app = new App({});
-  const stack = new Stack(app, 'test');
+  const stack = new Stack(app, name);
+
+  // This ensures that the 'aws:asset:path' metadata is set
+  stack.node.setContext(cxapi.ASSET_RESOURCE_METADATA_ENABLED_CONTEXT, true);
+
+  return { app, stack };
+}
+
+// Need to have CDK_OUTDIR set to something sensible as it's used to create the codePath when aws:asset:path is set
+const OLD_ENV = process.env;
+
+beforeEach(async () => {
+  jest.resetModules();
+  process.env = { ...OLD_ENV };
+  process.env.CDK_OUTDIR = await fs.mkdtemp(path.join(os.tmpdir(), 'uv-python-lambda-test-'));
+});
+
+afterEach(async () => {
+  if (process.env.CDK_OUTDIR) {
+    await fs.rm(process.env.CDK_OUTDIR, { recursive: true });
+  }
+  process.env = OLD_ENV;
+});
+
+test('Create a function from basic_app', async () => {
+  const { app, stack } = await createStack();
+
   new PythonFunction(stack, 'basic_app', {
     rootDir: path.join(resourcesPath, 'basic_app'),
     index: 'handler.py',
@@ -55,8 +89,8 @@ test('Create a function from basic_app', async () => {
 });
 
 test('Create a function from basic_app with no .py index extension', async () => {
-  const app = new App({});
-  const stack = new Stack(app, 'test');
+  const { stack } = await createStack();
+
   new PythonFunction(stack, 'basic_app', {
     rootDir: path.join(resourcesPath, 'basic_app'),
     index: 'handler',
@@ -77,9 +111,29 @@ test('Create a function from basic_app with no .py index extension', async () =>
   });
 });
 
+test('Create a function from basic_app when skip is true', async () => {
+  const { stack } = await createStack();
+
+  const bundlingSpy = jest.spyOn(stack, 'bundlingRequired', 'get').mockReturnValue(false);
+  const architecture = await getDockerHostArch();
+
+  // To see this fail, comment out the `if (skip) { return; } code in the PythonFunction constructor
+  expect(() => {
+    new PythonFunction(stack, 'basic_app', {
+      rootDir: path.join(resourcesPath, 'basic_app'),
+      index: 'handler',
+      handler: 'lambda_handler',
+      runtime: Runtime.PYTHON_3_12,
+      architecture,
+    });
+  }).not.toThrow();
+
+  bundlingSpy.mockRestore();
+});
+
 test('Create a function with workspaces_app', async () => {
-  const app = new App({});
-  const stack = new Stack(app, 'wstest');
+  const { app, stack } = await createStack('wstest');
+
   new PythonFunction(stack, 'workspaces_app', {
     rootDir: path.join(resourcesPath, 'workspaces_app'),
     workspacePackage: 'app',
@@ -122,4 +176,3 @@ async function getFunctionAssetContents(functionResource: any, app: App) {
   const contents = await fs.readdir(assetPath);
   return contents;
 }
-