From 8cbff710ce466134364b3f31d230f839365bde4b Mon Sep 17 00:00:00 2001
From: 0marperez <60363173+0marperez@users.noreply.github.com>
Date: Thu, 12 Dec 2024 09:05:17 -0500
Subject: [PATCH] misc: enable smoke tests (#1194)

---
 .../smoketests/SmokeTestsIntegration.kt       |  6 +++---
 .../smoketests/SmokeTestsRunnerGenerator.kt   | 21 +++++++++++--------
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/smoketests/SmokeTestsIntegration.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/smoketests/SmokeTestsIntegration.kt
index 9407c2911..192494e9e 100644
--- a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/smoketests/SmokeTestsIntegration.kt
+++ b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/smoketests/SmokeTestsIntegration.kt
@@ -5,6 +5,8 @@ import software.amazon.smithy.kotlin.codegen.core.CodegenContext
 import software.amazon.smithy.kotlin.codegen.core.DEFAULT_TEST_SOURCE_SET_ROOT
 import software.amazon.smithy.kotlin.codegen.core.KotlinDelegator
 import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.model.hasTrait
+import software.amazon.smithy.kotlin.codegen.utils.topDownOperations
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.smoketests.traits.SmokeTestsTrait
 
@@ -13,9 +15,7 @@ import software.amazon.smithy.smoketests.traits.SmokeTestsTrait
  */
 class SmokeTestsIntegration : KotlinIntegration {
     override fun enabledForService(model: Model, settings: KotlinSettings): Boolean =
-        false
-    // FIXME: Enable smoke tests again
-    // model.topDownOperations(settings.service).any { it.hasTrait<SmokeTestsTrait>() }
+        model.topDownOperations(settings.service).any { it.hasTrait<SmokeTestsTrait>() }
 
     override fun writeAdditionalFiles(ctx: CodegenContext, delegator: KotlinDelegator) =
         delegator.useFileWriter(
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/smoketests/SmokeTestsRunnerGenerator.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/smoketests/SmokeTestsRunnerGenerator.kt
index 81c275380..df8eeda53 100644
--- a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/smoketests/SmokeTestsRunnerGenerator.kt
+++ b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/smoketests/SmokeTestsRunnerGenerator.kt
@@ -30,6 +30,7 @@ object SmokeTestSectionIds {
     object HttpEngineOverride : SectionId
     object ServiceFilter : SectionId
     object SkipTags : SectionId
+    object SmokeTestsFile : SectionId
     object ClientConfig : SectionId {
         val Name: SectionKey<String> = SectionKey("aws.smithy.kotlin#SmokeTestClientConfigName")
         val Value: SectionKey<String> = SectionKey("aws.smithy.kotlin#SmokeTestClientConfigValue")
@@ -59,16 +60,18 @@ class SmokeTestsRunnerGenerator(
     ctx: CodegenContext,
 ) {
     internal fun render() {
-        writer.write("private var exitCode = 0")
-        renderEnvironmentVariables()
-        writer.declareSection(SmokeTestSectionIds.AdditionalEnvironmentVariables)
-        writer.write("")
-        writer.withBlock("public suspend fun main() {", "}") {
-            renderFunctionCalls()
-            write("#T(exitCode)", RuntimeTypes.Core.SmokeTests.exitProcess)
+        writer.declareSection(SmokeTestSectionIds.SmokeTestsFile) {
+            writer.write("private var exitCode = 0")
+            renderEnvironmentVariables()
+            writer.declareSection(SmokeTestSectionIds.AdditionalEnvironmentVariables)
+            writer.write("")
+            writer.withBlock("public suspend fun main() {", "}") {
+                renderFunctionCalls()
+                write("#T(exitCode)", RuntimeTypes.Core.SmokeTests.exitProcess)
+            }
+            writer.write("")
+            renderFunctions()
         }
-        writer.write("")
-        renderFunctions()
     }
 
     private fun renderEnvironmentVariables() {