From 5cd95591ce4fcf8e39164b3cd8ecba9f1bdc8c1b Mon Sep 17 00:00:00 2001
From: JordonPhillips <JordonPhillips@users.noreply.github.com>
Date: Tue, 26 Apr 2022 14:32:12 +0200
Subject: [PATCH] Make TestIndex tolerate synthetic traits

---
 .../smithy/model/knowledge/TextIndex.java     |  7 ++--
 .../smithy/model/knowledge/TextIndexTest.java | 37 +++++++++++++++++++
 .../smithy/model/knowledge/text-index.smithy  |  8 ++++
 3 files changed, 49 insertions(+), 3 deletions(-)
 create mode 100644 smithy-model/src/test/java/software/amazon/smithy/model/knowledge/TextIndexTest.java
 create mode 100644 smithy-model/src/test/resources/software/amazon/smithy/model/knowledge/text-index.smithy

diff --git a/smithy-model/src/main/java/software/amazon/smithy/model/knowledge/TextIndex.java b/smithy-model/src/main/java/software/amazon/smithy/model/knowledge/TextIndex.java
index 090788b853f..3525f2bf139 100644
--- a/smithy-model/src/main/java/software/amazon/smithy/model/knowledge/TextIndex.java
+++ b/smithy-model/src/main/java/software/amazon/smithy/model/knowledge/TextIndex.java
@@ -75,9 +75,10 @@ private static void computeShapeTextInstances(
         textInstances.add(TextInstance.createShapeInstance(shape));
 
         for (Trait trait : shape.getAllTraits().values()) {
-            Shape traitShape = model.expectShape(trait.toShapeId());
-            computeTextInstancesForAppliedTrait(trait.toNode(), trait, shape, textInstances,
-                    new ArrayDeque<>(), model, traitShape);
+            model.getShape(trait.toShapeId()).ifPresent(traitShape -> {
+                computeTextInstancesForAppliedTrait(trait.toNode(), trait, shape, textInstances,
+                        new ArrayDeque<>(), model, traitShape);
+            });
         }
     }
 
diff --git a/smithy-model/src/test/java/software/amazon/smithy/model/knowledge/TextIndexTest.java b/smithy-model/src/test/java/software/amazon/smithy/model/knowledge/TextIndexTest.java
new file mode 100644
index 00000000000..08544d33e14
--- /dev/null
+++ b/smithy-model/src/test/java/software/amazon/smithy/model/knowledge/TextIndexTest.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *   http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.smithy.model.knowledge;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.empty;
+import static org.hamcrest.Matchers.hasSize;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.smithy.model.Model;
+
+public class TextIndexTest {
+
+    @Test
+    public void handlesSyntheticTraits() {
+        Model model = Model.assembler()
+                .addImport(OperationIndexTest.class.getResource("text-index.smithy"))
+                .assemble()
+                .unwrap();
+
+        TextIndex index = TextIndex.of(model);
+        assertThat(index.getTextInstances(), hasSize(5));
+    }
+}
diff --git a/smithy-model/src/test/resources/software/amazon/smithy/model/knowledge/text-index.smithy b/smithy-model/src/test/resources/software/amazon/smithy/model/knowledge/text-index.smithy
new file mode 100644
index 00000000000..6a2f8d54120
--- /dev/null
+++ b/smithy-model/src/test/resources/software/amazon/smithy/model/knowledge/text-index.smithy
@@ -0,0 +1,8 @@
+$version: "2.0"
+
+namespace com.example
+
+enum FooEnum {
+    /// Documentation
+    FOO
+}