From 8eee5876d0ac163dbe1e3588c80956f3fcada574 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olivier=20M=C3=A9lois?= <olivier.melois@disneystreaming.com>
Date: Thu, 30 May 2024 12:55:23 +0200
Subject: [PATCH] Fix TraitProviderSpec

---
 modules/core/test/src/TraitProviderSpec.scala | 34 ++++++++++++-------
 1 file changed, 22 insertions(+), 12 deletions(-)

diff --git a/modules/core/test/src/TraitProviderSpec.scala b/modules/core/test/src/TraitProviderSpec.scala
index 6344c00..16f8cc3 100644
--- a/modules/core/test/src/TraitProviderSpec.scala
+++ b/modules/core/test/src/TraitProviderSpec.scala
@@ -19,6 +19,7 @@ import munit.Location
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.traits.Trait
 import software.amazon.smithy.model.validation.ValidatedResult
+import scala.jdk.CollectionConverters._
 
 import scala.io.Source
 
@@ -37,19 +38,28 @@ final class SanitySpec extends munit.FunSuite {
       )
     }
 
-    val lines = scala.util
-      .Using(
-        Source.fromResource(
-          "META-INF/services/software.amazon.smithy.model.traits.TraitService"
-        )
-      ) {
-        _.getLines().toList
-      }
-      .fold(
-        ex => fail("Failed to load TraitService resource", ex),
-        identity
+    val resources = this
+      .getClass()
+      .getClassLoader()
+      .getResources(
+        "META-INF/services/software.amazon.smithy.model.traits.TraitService"
       )
-    val classesFQN = lines.map(_.split("\\$").head)
+      .asIterator()
+      .asScala
+      .toList
+
+    val lines = resources.flatMap { resource =>
+      scala.util
+        .Using(Source.fromURL(resource)) {
+          _.getLines().toList
+        }
+        .fold(
+          ex => fail("Failed to load TraitService resource", ex),
+          identity
+        )
+    }
+    val classesFQN =
+      lines.map(_.split("\\$").head).filter(_.startsWith("alloy"))
     val classes =
       classesFQN.map(name => this.getClass().getClassLoader().loadClass(name))
     assert(