Add ModelTransformerPlugin for endpoint test cases

This commit adds a ModelTransformerPlugin implementation that removes
test case entries that rely on operations being removed from a model.
If all test cases are removed, the trait is also removed.

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/traits/CleanEndpointTestOperationInput.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rulesengine.traits;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.OperationShape;
+import software.amazon.smithy.model.shapes.ServiceShape;
+import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.transform.ModelTransformer;
+import software.amazon.smithy.model.transform.ModelTransformerPlugin;
+import software.amazon.smithy.utils.SmithyInternalApi;
+
+@SmithyInternalApi
+public class CleanEndpointTestOperationInput implements ModelTransformerPlugin {
+    @Override
+    public Model onRemove(ModelTransformer transformer, Collection<Shape> removed, Model model) {
+        Set<Shape> servicesToUpdate = getServicesToUpdate(model, removed);
+        return transformer.replaceShapes(model, servicesToUpdate);
+    }
+
+    private Set<Shape> getServicesToUpdate(Model model, Collection<Shape> removed) {
+        // Precompute shape ids to operations that were removed.
+        Map<String, OperationShape> removedOperationMap = new HashMap<>();
+        for (Shape shape : removed) {
+            if (shape.isOperationShape()) {
+                removedOperationMap.put(shape.getId().getName(), shape.asOperationShape().get());
+            }
+        }
+
+        Set<Shape> result = new HashSet<>();
+        for (ServiceShape serviceShape : model.getServiceShapesWithTrait(EndpointTestsTrait.class)) {
+            EndpointTestsTrait trait = serviceShape.expectTrait(EndpointTestsTrait.class);
+            List<EndpointTestCase> updatedTestCases = new ArrayList<>(trait.getTestCases());
+            // Check each input to each test case and remove entries from the list.
+            for (EndpointTestCase testCase : trait.getTestCases()) {
+                for (EndpointTestOperationInput input : testCase.getOperationInputs()) {
+                    if (removed.contains(removedOperationMap.get(input.getOperationName()))) {
+                        updatedTestCases.remove(testCase);
+                    }
+                }
+            }
+
+            // Update the shape if the trait has changed, removing if no cases are left.
+            if (updatedTestCases.isEmpty()) {
+                result.add(serviceShape.toBuilder().removeTrait(EndpointTestsTrait.ID).build());
+            } else if (updatedTestCases.size() != trait.getTestCases().size()) {
+                result.add(serviceShape.toBuilder()
+                        .addTrait(trait.toBuilder().testCases(updatedTestCases).build())
+                        .build());
+            }
+        }
+
+        return result;
+    }
+}

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/traits/ContextParamTrait.java

@@ -46,7 +46,7 @@ protected Node createNode() {
     }
 
     @Override
-    public SmithyBuilder<ContextParamTrait> toBuilder() {
+    public Builder toBuilder() {
         return new Builder()
                 .sourceLocation(getSourceLocation())
                 .name(name);

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/traits/EndpointRuleSetTrait.java

@@ -53,7 +53,7 @@ protected Node createNode() {
     }
 
     @Override
-    public SmithyBuilder<EndpointRuleSetTrait> toBuilder() {
+    public Builder toBuilder() {
         return builder()
                 .sourceLocation(getSourceLocation())
                 .ruleSet(ruleSet);

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/traits/EndpointTestsTrait.java

@@ -61,7 +61,7 @@ protected Node createNode() {
     }
 
     @Override
-    public SmithyBuilder<EndpointTestsTrait> toBuilder() {
+    public Builder toBuilder() {
         return builder()
                 .sourceLocation(getSourceLocation())
                 .version(version)

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/traits/EndpointTestsTraitValidator.java

@@ -47,9 +47,12 @@ public List<ValidationEvent> validate(Model model) {
                                         operationName, serviceShape.getId())));
                     }
 
-                    StructureShape inputShape = model.expectShape(
-                            operationNameMap.get(operationName).getInputShape(), StructureShape.class);
-                    events.addAll(validateOperationInput(model, serviceShape, inputShape, testOperationInput));
+                    // Still emit events if the operation exists, but was just not bound.
+                    if (operationNameMap.containsKey(operationName)) {
+                        StructureShape inputShape = model.expectShape(
+                                operationNameMap.get(operationName).getInputShape(), StructureShape.class);
+                        events.addAll(validateOperationInput(model, serviceShape, inputShape, testOperationInput));
+                    }
                 }
             }
         }

smithy-rules-engine/src/main/java/software/amazon/smithy/rulesengine/traits/StaticContextParamsTrait.java

@@ -14,7 +14,6 @@
 import software.amazon.smithy.model.traits.AbstractTraitBuilder;
 import software.amazon.smithy.model.traits.Trait;
 import software.amazon.smithy.utils.BuilderRef;
-import software.amazon.smithy.utils.SmithyBuilder;
 import software.amazon.smithy.utils.SmithyUnstableApi;
 import software.amazon.smithy.utils.ToSmithyBuilder;
 
@@ -48,7 +47,7 @@ protected Node createNode() {
     }
 
     @Override
-    public SmithyBuilder<StaticContextParamsTrait> toBuilder() {
+    public Builder toBuilder() {
         return new Builder()
                 .sourceLocation(getSourceLocation())
                 .parameters(parameters);

smithy-rules-engine/src/main/resources/META-INF/services/software.amazon.smithy.model.transform.ModelTransformerPlugin

@@ -0,0 +1 @@
+software.amazon.smithy.rulesengine.traits.CleanEndpointTestOperationInput

smithy-rules-engine/src/test/java/software/amazon/smithy/rulesengine/traits/CleanEndpointTestOperationInputTest.java

@@ -0,0 +1,60 @@
+package software.amazon.smithy.rulesengine.traits;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.shapes.ServiceShape;
+import software.amazon.smithy.model.shapes.ShapeId;
+import software.amazon.smithy.model.transform.ModelTransformer;
+
+public class CleanEndpointTestOperationInputTest {
+    private static final ShapeId SERVICE_ID = ShapeId.from("smithy.example#ExampleService");
+    private static final ShapeId GET_THING = ShapeId.from("smithy.example#GetThing");
+    private static final ShapeId PING = ShapeId.from("smithy.example#Ping");
+    private static Model model;
+
+    @BeforeAll
+    public static void before() {
+        model = Model.assembler()
+                .discoverModels(ContextIndexTest.class.getClassLoader())
+                .addImport(ContextIndexTest.class.getResource("traits-test-model.smithy"))
+                .assemble()
+                .unwrap();
+    }
+
+    @Test
+    public void retainsTestsIfOperationRemains() {
+        Model transformed = ModelTransformer.create().filterShapes(model, shape -> !shape.getId().equals(PING));
+
+        assertFalse(transformed.getShape(PING).isPresent());
+        assertTrue(transformed.getShape(SERVICE_ID).isPresent());
+
+        ServiceShape mainService = model.expectShape(SERVICE_ID, ServiceShape.class);
+        assertTrue(mainService.hasTrait(EndpointTestsTrait.class));
+        ServiceShape transformedService = transformed.expectShape(SERVICE_ID, ServiceShape.class);
+        assertTrue(transformedService.hasTrait(EndpointTestsTrait.class));
+
+        Node.assertEquals(transformedService.expectTrait(EndpointTestsTrait.class).toNode(),
+                mainService.expectTrait(EndpointTestsTrait.class).toNode());
+    }
+
+    @Test
+    public void removesTestsIfOperationRemoved() {
+        Model transformed = ModelTransformer.create().filterShapes(model, shape -> !shape.getId().equals(GET_THING));
+
+        assertFalse(transformed.getShape(GET_THING).isPresent());
+        assertTrue(transformed.getShape(SERVICE_ID).isPresent());
+
+        ServiceShape transformedService = transformed.expectShape(SERVICE_ID, ServiceShape.class);
+        assertTrue(transformedService.hasTrait(EndpointTestsTrait.class));
+
+        EndpointTestsTrait trait = transformedService.expectTrait(EndpointTestsTrait.class);
+        assertEquals(1, trait.getTestCases().size());
+        assertTrue(trait.getTestCases().get(0).getOperationInputs().isEmpty());
+    }
+}

smithy-rules-engine/src/test/resources/software/amazon/smithy/rulesengine/traits/traits-test-model.smithy

@@ -15,7 +15,7 @@ use smithy.rules#staticContextParams
 @suppress(["RuleSetParameter.TestCase.Unused"])
 service ExampleService {
     version: "2022-01-01",
-    operations: [GetThing]
+    operations: [GetThing, Ping]
 }
 
 apply ExampleService @endpointRuleSet({
@@ -152,3 +152,5 @@ structure GetThingInput {
     @contextParam(name: "boolBaz")
     fuzz: String,
 }
+
+operation Ping {}