Merge pull request #688 from jlbutler/fix-union-original-shape-name

Fix OriginalShapeName handling for union types

Author: knottnt

pkg/api/passes_test.go

@@ -974,3 +974,91 @@ func TestValidateShapeNameMethod(t *testing.T) {
 		})
 	}
 }
+
+// TestRenamedUnionShapePreservesOriginalName verifies that when a union
+// shape is renamed (e.g. MemoryStrategyInput -> MemoryStrategyInput_) by
+// renameIOSuffixedShapeNames, the OriginalShapeName is preserved. The
+// downstream SDK code generation functions (setSDKForUnion,
+// setResourceForUnion, varEmptyConstructorSDKType) use OriginalShapeName
+// to emit correct SDK type references.
+func TestRenamedUnionShapePreservesOriginalName(t *testing.T) {
+	a := &API{
+		name: "testapi",
+		Metadata: Metadata{
+			APIVersion:          "0000-00-00",
+			EndpointPrefix:      "testapi",
+			JSONVersion:         "1.1",
+			Protocol:            "json",
+			ServiceAbbreviation: "TestAPI",
+			ServiceFullName:     "Test API",
+			SignatureVersion:    "v4",
+		},
+		Operations: map[string]*Operation{
+			"CreateMemory": {
+				Name:      "CreateMemory",
+				InputRef:  ShapeRef{ShapeName: "CreateMemoryInput"},
+				OutputRef: ShapeRef{ShapeName: "CreateMemoryOutput"},
+			},
+		},
+		Shapes: map[string]*Shape{
+			"CreateMemoryInput": {
+				ShapeName: "CreateMemoryInput",
+				Type:      "structure",
+				MemberRefs: map[string]*ShapeRef{
+					"Strategies": {ShapeName: "MemoryStrategyInputList"},
+				},
+			},
+			"CreateMemoryOutput": {
+				ShapeName: "CreateMemoryOutput",
+				Type:      "structure",
+			},
+			"MemoryStrategyInputList": {
+				ShapeName: "MemoryStrategyInputList",
+				Type:      "list",
+				MemberRef: ShapeRef{ShapeName: "MemoryStrategyInput"},
+			},
+			"MemoryStrategyInput": {
+				ShapeName: "MemoryStrategyInput",
+				Type:      "structure",
+				RealType:  "union",
+			},
+		},
+	}
+
+	// Wire up shape refs
+	for _, op := range a.Operations {
+		op.InputRef.API = a
+		op.InputRef.Shape = a.Shapes[op.InputRef.ShapeName]
+		op.OutputRef.API = a
+		op.OutputRef.Shape = a.Shapes[op.OutputRef.ShapeName]
+	}
+	for _, s := range a.Shapes {
+		s.API = a
+		for k := range s.MemberRefs {
+			ref := s.MemberRefs[k]
+			ref.API = a
+			ref.Shape = a.Shapes[ref.ShapeName]
+			s.MemberRefs[k] = ref
+		}
+		if s.MemberRef.ShapeName != "" {
+			s.MemberRef.API = a
+			s.MemberRef.Shape = a.Shapes[s.MemberRef.ShapeName]
+		}
+	}
+
+	unionShape := a.Shapes["MemoryStrategyInput"]
+	if unionShape.OriginalShapeName != "" {
+		t.Fatalf("expected empty OriginalShapeName before rename, got %q", unionShape.OriginalShapeName)
+	}
+
+	a.renameIOSuffixedShapeNames()
+
+	// ShapeName gets underscore suffix, OriginalShapeName preserves the
+	// pre-rename value for downstream SDK code generation.
+	if unionShape.ShapeName != "MemoryStrategyInput_" {
+		t.Fatalf("expected ShapeName %q after rename, got %q", "MemoryStrategyInput_", unionShape.ShapeName)
+	}
+	if unionShape.OriginalShapeName != "MemoryStrategyInput" {
+		t.Fatalf("expected OriginalShapeName %q after rename, got %q", "MemoryStrategyInput", unionShape.OriginalShapeName)
+	}
+}

pkg/generate/code/set_resource.go

@@ -2454,6 +2454,11 @@ func setResourceForUnion(
 
 	sdkGoType := sourceShape.GoTypeWithPkgName()
 	sdkGoType = model.ReplacePkgName(sdkGoType, r.SDKAPIPackageName(), "svcsdktypes", true)
+	// Use the original shape name for SDK type references when the shape
+	// was renamed (e.g. Input/Output suffix collision avoidance).
+	if sourceShape.OriginalShapeName != "" {
+		sdkGoType = "*svcsdktypes." + sourceShape.OriginalShapeName
+	}
 
 	out += fmt.Sprintf("%sswitch %s.(type) {\n", indent, sourceVarName)
 	for _, targetMemberName := range targetShape.MemberNames() {

pkg/generate/code/set_resource_test.go

@@ -5579,3 +5579,31 @@ func TestSetResource_BedrockAgentCoreControl_GatewayTarget_NestedUnionTypeSwitch
 	assert.Contains(got, "case *svcsdktypes.ApiSchemaConfigurationMemberInlinePayload:")
 	assert.Contains(got, "case *svcsdktypes.ApiSchemaConfigurationMemberS3:")
 }
+
+// TestSetResource_BedrockAgentCoreControl_Memory_InputSuffixUnion tests that
+// union types whose names end in "Input" (e.g. MemoryStrategyInput) use the
+// original SDK shape name in generated set-output code, not the renamed
+// version with an underscore suffix (MemoryStrategyInput_).
+func TestSetResource_BedrockAgentCoreControl_Memory_InputSuffixUnion(t *testing.T) {
+	assert := assert.New(t)
+	require := require.New(t)
+
+	g := testutil.NewModelForServiceWithOptions(t, "bedrock-agentcore-control", &testutil.TestingModelOptions{
+		GeneratorConfigFile: "generator-with-input-suffix-union.yaml",
+	})
+
+	crd := testutil.GetCRDByName(t, g, "Memory")
+	require.NotNil(crd)
+	assert.NotNil(crd.Ops.ReadOne)
+
+	got, err := code.SetResource(crd.Config(), crd, model.OpTypeGet, "resp", "ko", 1)
+	require.NoError(err)
+
+	// On the read path, union type switches must use the original SDK name
+	assert.NotContains(got, "MemoryStrategyInput_",
+		"Should use original SDK shape name MemoryStrategyInput, not renamed MemoryStrategyInput_")
+	assert.NotContains(got, "CustomConfigurationInput_",
+		"Should use original SDK shape name CustomConfigurationInput, not renamed version")
+	assert.NotContains(got, "TriggerConditionInput_",
+		"Should use original SDK shape name TriggerConditionInput, not renamed version")
+}

pkg/generate/code/set_sdk.go

@@ -1533,7 +1533,7 @@ func varEmptyConstructorSDKType(
 	// For SDK types, we need to use the original shape name (before stutter
 	// removal) since the AWS SDK uses the original names. The stutter removal
 	// renames are only for CRD types.
-	if shape.Type == "structure" && shape.OriginalShapeName != "" {
+	if (shape.Type == "structure" || shape.RealType == "union") && shape.OriginalShapeName != "" {
 		// Replace the renamed shape name with the original SDK shape name
 		goType = "svcsdktypes." + shape.OriginalShapeName
 	}
@@ -1863,6 +1863,11 @@ func setSDKForUnion(
 
 	sdkGoType := targetShape.GoTypeWithPkgName()
 	sdkGoType = model.ReplacePkgName(sdkGoType, r.SDKAPIPackageName(), "svcsdktypes", false)
+	// Use the original shape name for SDK type references when the shape
+	// was renamed (e.g. Input/Output suffix collision avoidance).
+	if targetShape.OriginalShapeName != "" {
+		sdkGoType = "svcsdktypes." + targetShape.OriginalShapeName
+	}
 
 	out += fmt.Sprintf("%sisInterfaceSet := false\n", indent)
 

pkg/generate/code/set_sdk_test.go

@@ -6769,3 +6769,40 @@ func TestSetSDK_BedrockAgentCoreControl_GatewayTarget_NestedUnionNoDereference(t
 	// dereferenced (e.g., ApiGatewayTargetConfiguration is a regular struct)
 	assert.Contains(got, "f7f0f0Parent.Value = *f7f0f0\n")
 }
+
+// TestSetSDK_BedrockAgentCoreControl_Memory_InputSuffixUnion tests that
+// union types whose names end in "Input" (e.g. MemoryStrategyInput) use
+// the original SDK shape name in generated code, not the renamed version
+// with an underscore suffix (MemoryStrategyInput_). The rename pass adds
+// the underscore to avoid collisions with operation input types, but SDK
+// type references must use the original name.
+func TestSetSDK_BedrockAgentCoreControl_Memory_InputSuffixUnion(t *testing.T) {
+	assert := assert.New(t)
+	require := require.New(t)
+
+	g := testutil.NewModelForServiceWithOptions(t, "bedrock-agentcore-control", &testutil.TestingModelOptions{
+		GeneratorConfigFile: "generator-with-input-suffix-union.yaml",
+	})
+
+	crd := testutil.GetCRDByName(t, g, "Memory")
+	require.NotNil(crd)
+	assert.NotNil(crd.Ops.Create)
+
+	got, err := code.SetSDK(crd.Config(), crd, model.OpTypeCreate, "r.ko", "res", 1)
+	require.NoError(err)
+
+	// MemoryStrategyInput is a union whose name ends in "Input". The code-gen
+	// renames it to MemoryStrategyInput_ to avoid collision with operation
+	// input types. SDK references must use the original name.
+	assert.Contains(got, "svcsdktypes.MemoryStrategyInput")
+	assert.NotContains(got, "MemoryStrategyInput_",
+		"Should use original SDK shape name MemoryStrategyInput, not renamed MemoryStrategyInput_")
+
+	// CustomConfigurationInput is also a union ending in "Input"
+	assert.NotContains(got, "CustomConfigurationInput_",
+		"Should use original SDK shape name CustomConfigurationInput, not renamed version")
+
+	// TriggerConditionInput is also a union ending in "Input"
+	assert.NotContains(got, "TriggerConditionInput_",
+		"Should use original SDK shape name TriggerConditionInput, not renamed version")
+}

pkg/testdata/models/apis/bedrock-agentcore-control/0000-00-00/generator-with-input-suffix-union.yaml

@@ -0,0 +1,33 @@
+ignore:
+  resource_names:
+    - AgentRuntime
+    - AgentRuntimeEndpoint
+    - ApiKeyCredentialProvider
+    - Browser
+    - BrowserProfile
+    - CodeInterpreter
+    - Evaluator
+    - Gateway
+    - GatewayTarget
+    # Memory
+    - Oauth2CredentialProvider
+    - OnlineEvaluationConfig
+    - WorkloadIdentity
+    - Policy
+    - PolicyEngine
+  field_paths:
+    - MemoryStrategy.Configuration
+    - CreateMemoryInput.ClientToken
+    - UpdateMemoryInput.ClientToken
+    - DeleteMemoryInput.ClientToken
+sdk_names:
+  model_name: bedrock-agentcore-control
+
+resources:
+  Memory:
+    ignore_idempotency_token: true
+    fields:
+      MemoryId:
+        is_primary_key: true
+    tags:
+      ignore: true