Repository: pulumi/pulumi-awsx
Pull request #1354 (2024-08-01 22:55:38)

refactor: clarify how Vpc subnet allocators are chosen
The case analysis around picking which allocator to use is made more explicit and covered with a unit test.


Patch:
diff --git a/awsx/ec2/subnetDistributorNew.ts b/awsx/ec2/subnetDistributorNew.ts
index 7224ce3b..8c55fbec 100644
--- a/awsx/ec2/subnetDistributorNew.ts
+++ b/awsx/ec2/subnetDistributorNew.ts
@@ -169,16 +169,18 @@ function subnetName(vpcName: string, subnet: SubnetSpecInputs, azNum: number): s
   return `${vpcName}-${specName}-${azNum}`;
 }
 
+export type NormalizedSubnetInputs =
+  | { normalizedSpecs: SubnetSpecInputs[]; isExplicitLayout: false }
+  | { normalizedSpecs: ExplicitSubnetSpecInputs[]; isExplicitLayout: true }
+  | undefined;
+
 /* Ensure all inputs are consistent and fill in missing values with defaults
  * Ensure any specified, netmask, size or blocks are in agreement.
  */
 export function validateAndNormalizeSubnetInputs(
   subnetArgs: SubnetSpecInputs[] | undefined,
   availabilityZoneCount: number,
-):
-  | { normalizedSpecs: SubnetSpecInputs[]; isExplicitLayout: false }
-  | { normalizedSpecs: ExplicitSubnetSpecInputs[]; isExplicitLayout: true }
-  | undefined {
+): NormalizedSubnetInputs {
   if (subnetArgs === undefined) {
     return undefined;
   }
diff --git a/awsx/ec2/vpc.ts b/awsx/ec2/vpc.ts
index 14eb9543..322d9ad7 100644
--- a/awsx/ec2/vpc.ts
+++ b/awsx/ec2/vpc.ts
@@ -20,6 +20,8 @@ import {
   getSubnetSpecs,
   getSubnetSpecsExplicit,
   validateAndNormalizeSubnetInputs,
+  NormalizedSubnetInputs,
+  ExplicitSubnetSpecInputs,
 } from "./subnetDistributorNew";
 import { Netmask } from "netmask";
 
@@ -87,33 +89,35 @@ export class Vpc extends schema.Vpc<VpcData> {
 
     const cidrBlock = args.cidrBlock ?? "10.0.0.0/16";
 
-    const parsedSpecs = validateAndNormalizeSubnetInputs(
+    const parsedSpecs: NormalizedSubnetInputs = validateAndNormalizeSubnetInputs(
       args.subnetSpecs,
       availabilityZones.length,
     );
 
     const subnetStrategy = args.subnetStrategy ?? "Legacy";
     const subnetSpecs = (() => {
-      if (parsedSpecs === undefined || subnetStrategy === "Legacy") {
-        const legacySubnetSpecs = getSubnetSpecsLegacy(
-          name,
-          cidrBlock,
-          availabilityZones,
-          parsedSpecs?.normalizedSpecs,
-        );
-        return legacySubnetSpecs;
-      }
-
-      if (parsedSpecs.isExplicitLayout) {
-        return getSubnetSpecsExplicit(name, availabilityZones, parsedSpecs.normalizedSpecs);
+      const a = Vpc.pickSubnetAllocator(parsedSpecs, subnetStrategy);
+      switch (a.allocator) {
+        case "LegacyAllocator":
+          const legacySubnetSpecs = getSubnetSpecsLegacy(
+            name,
+            cidrBlock,
+            availabilityZones,
+            parsedSpecs?.normalizedSpecs,
+          );
+          return legacySubnetSpecs;
+        case "ExplicitAllocator":
+          return getSubnetSpecsExplicit(name, availabilityZones, a.specs);
+        case "NewAllocator":
+        default:
+          return getSubnetSpecs(
+            name,
+            cidrBlock,
+            availabilityZones,
+            parsedSpecs?.normalizedSpecs,
+            args.availabilityZoneCidrMask,
+          );
       }
-      return getSubnetSpecs(
-        name,
-        cidrBlock,
-        availabilityZones,
-        parsedSpecs.normalizedSpecs,
-        args.availabilityZoneCidrMask,
-      );
     })();
 
     let subnetLayout = parsedSpecs?.normalizedSpecs;
@@ -361,6 +365,22 @@ export class Vpc extends schema.Vpc<VpcData> {
     }
     return result.names.slice(0, desiredCount);
   }
+
+  // Internal. Exported for testing.
+  public static pickSubnetAllocator(
+    parsedSpecs: NormalizedSubnetInputs,
+    subnetStrategy: schema.SubnetAllocationStrategyInputs,
+  ):
+    | { allocator: "LegacyAllocator" | "NewAllocator" }
+    | { allocator: "ExplicitAllocator"; specs: ExplicitSubnetSpecInputs[] } {
+    if (parsedSpecs === undefined || subnetStrategy === "Legacy") {
+      return { allocator: "LegacyAllocator" };
+    }
+    if (parsedSpecs.isExplicitLayout) {
+      return { allocator: "ExplicitAllocator", specs: parsedSpecs.normalizedSpecs };
+    }
+    return { allocator: "NewAllocator" };
+  }
 }
 
 export function extractSubnetSpecInputFromLegacyLayout(


Test patch:
diff --git a/awsx/ec2/vpc.test.ts b/awsx/ec2/vpc.test.ts
index f9c17aa2..9f859fba 100644
--- a/awsx/ec2/vpc.test.ts
+++ b/awsx/ec2/vpc.test.ts
@@ -14,7 +14,11 @@
 
 import fc from "fast-check";
 import { getOverlappingSubnets } from ".";
-import { NatGatewayStrategyInputs, SubnetTypeInputs } from "../schema-types";
+import {
+  NatGatewayStrategyInputs,
+  SubnetTypeInputs,
+  SubnetAllocationStrategyInputs,
+} from "../schema-types";
 import { getSubnetSpecsLegacy } from "./subnetDistributorLegacy";
 import {
   compareSubnetSpecs,
@@ -24,6 +28,7 @@ import {
   validateEips,
   validateNatGatewayStrategy,
   validateSubnets,
+  Vpc,
 } from "./vpc";
 import { Netmask, long2ip, ip2long } from "netmask";
 
@@ -291,3 +296,50 @@ describe("finding subnet gaps", () => {
     );
   });
 });
+
+describe("picking subnet allocator", () => {
+  it("picks legacy allocator for the legacy strategy", () => {
+    const a = Vpc.pickSubnetAllocator(undefined, "Legacy");
+    expect(a.allocator).toBe("LegacyAllocator");
+  });
+
+  // This behavior is suspect as NewAllocator supports working without any subnet specs defined.
+  it("picks legacy allocator for the auto strategy with no specs", () => {
+    const a = Vpc.pickSubnetAllocator(undefined, "Auto");
+    expect(a.allocator).toBe("LegacyAllocator");
+  });
+
+  // This behavior looks like a degenerate case perhaps marking the strategy as invalid and failing could be preferable.
+  it("picks legacy allocator for the exact strategy with no specs", () => {
+    const a = Vpc.pickSubnetAllocator(undefined, "Exact");
+    expect(a.allocator).toBe("LegacyAllocator");
+  });
+
+  it("picks explicit allocator for explicit layout unless legacy is asked for", () => {
+    const f = (strat: SubnetAllocationStrategyInputs) =>
+      Vpc.pickSubnetAllocator(
+        {
+          normalizedSpecs: [{ type: "Private", cidrBlocks: ["10.2.0.0/16"] }],
+          isExplicitLayout: true,
+        },
+        strat,
+      ).allocator;
+    expect(f("Legacy")).toBe("LegacyAllocator");
+    expect(f("Auto")).toBe("ExplicitAllocator");
+    expect(f("Exact")).toBe("ExplicitAllocator");
+  });
+
+  it("picks new allocator for non-explicit layout unless legacy is asked for", () => {
+    const f = (strat: SubnetAllocationStrategyInputs) =>
+      Vpc.pickSubnetAllocator(
+        {
+          normalizedSpecs: [{ type: "Private" }],
+          isExplicitLayout: false,
+        },
+        strat,
+      ).allocator;
+    expect(f("Legacy")).toBe("LegacyAllocator");
+    expect(f("Auto")).toBe("NewAllocator");
+    expect(f("Exact")).toBe("NewAllocator"); // this case is a bit suspect
+  });
+});