Retain endpointTests trait if all cases removed

smithy-lang · Sep 20, 2023 · d8aa29d · d8aa29d
1 parent 5ee9c18
commit d8aa29d
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 4 deletions.
diff --git a/.../main/java/software/amazon/smithy/rulesengine/traits/CleanEndpointTestOperationInput.java b/.../main/java/software/amazon/smithy/rulesengine/traits/CleanEndpointTestOperationInput.java
@@ -40,6 +40,10 @@ private Set<Shape> getServicesToUpdate(Model model, Collection<Shape> removed) {
         Set<Shape> result = new HashSet<>();
         for (ServiceShape serviceShape : model.getServiceShapesWithTrait(EndpointTestsTrait.class)) {
             EndpointTestsTrait trait = serviceShape.expectTrait(EndpointTestsTrait.class);
+            if (trait.getTestCases().isEmpty()) {
+                continue;
+            }
+
             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()) {
@@ -50,10 +54,8 @@ private Set<Shape> getServicesToUpdate(Model model, Collection<Shape> removed) {
                 }
             }
 
-            // 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()) {
+            // Update the shape if the trait has changed.
+            if (updatedTestCases.size() != trait.getTestCases().size()) {
                 result.add(serviceShape.toBuilder()
                         .addTrait(trait.toBuilder().testCases(updatedTestCases).build())
                         .build());

diff --git a/...t/java/software/amazon/smithy/rulesengine/traits/CleanEndpointTestOperationInputTest.java b/...t/java/software/amazon/smithy/rulesengine/traits/CleanEndpointTestOperationInputTest.java
@@ -11,6 +11,7 @@
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.transform.ModelTransformer;
+import software.amazon.smithy.utils.ListUtils;
 
 public class CleanEndpointTestOperationInputTest {
     private static final ShapeId SERVICE_ID = ShapeId.from("smithy.example#ExampleService");
@@ -57,4 +58,28 @@ public void removesTestsIfOperationRemoved() {
         assertEquals(1, trait.getTestCases().size());
         assertTrue(trait.getTestCases().get(0).getOperationInputs().isEmpty());
     }
+
+    @Test
+    public void retainsTraitIfAllTestsRemoved() {
+        ServiceShape serviceShape = model.expectShape(SERVICE_ID, ServiceShape.class);
+        EndpointTestsTrait replacementTrait = serviceShape.expectTrait(EndpointTestsTrait.class);
+
+        // Hack out the test case without operation input.
+        ModelTransformer modelTransformer = ModelTransformer.create();
+        replacementTrait = replacementTrait.toBuilder().removeTestCase(replacementTrait.getTestCases().get(0)).build();
+        Model transformed = modelTransformer.replaceShapes(model, ListUtils.of(
+                serviceShape.toBuilder().addTrait(replacementTrait).build()));
+
+        // Then do the filtering.
+        transformed = modelTransformer.filterShapes(transformed, 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(0, trait.getTestCases().size());
+    }
 }