Defer file writer creation for protocol tests

Author: JordonPhillips

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/CodegenVisitor.java

@@ -225,10 +225,9 @@ void execute() {
             String baseName = protocol.getName().toLowerCase(Locale.US)
                     .replace("-", "_")
                     .replace(".", "_");
-            writers.useFileWriter(String.format("tests/functional/%s.spec.ts", baseName), writer -> {
-                context.setWriter(writer);
-                protocolGenerator.generateProtocolTests(context);
-            });
+            String protocolTestFileName = String.format("tests/functional/%s.spec.ts", baseName);
+            context.setDeferredWriter(() -> writers.checkoutFileWriter(protocolTestFileName));
+            protocolGenerator.generateProtocolTests(context);
         }
 
         // Write each pending writer.

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/HttpProtocolTestGenerator.java

@@ -96,9 +96,9 @@ public final class HttpProtocolTestGenerator implements Runnable {
     private final Set<String> additionalStubs = new TreeSet<>();
     private final ProtocolGenerator protocolGenerator;
     private final TestFilter testFilter;
+    private final GenerationContext context;
 
     private TypeScriptWriter writer;
-    private boolean writerInitialized = false;
 
     HttpProtocolTestGenerator(
             GenerationContext context,
@@ -110,10 +110,10 @@ public final class HttpProtocolTestGenerator implements Runnable {
         this.protocol = context.getSettings().getProtocol();
         this.service = settings.getService(model);
         this.symbolProvider = context.getSymbolProvider();
-        this.writer = context.getWriter();
         this.protocolGenerator = protocolGenerator;
         serviceSymbol = symbolProvider.toSymbol(service);
         this.testFilter = testFilter;
+        this.context = context;
     }
 
     HttpProtocolTestGenerator(
@@ -210,12 +210,12 @@ private <T extends HttpMessageTestCase> void onlyIfProtocolMatches(T testCase, R
     }
 
     private void initializeWriterIfNeeded() {
-        if (!writerInitialized) {
+        if (writer == null) {
+            writer = context.getWriter();
             writer.addDependency(TypeScriptDependency.AWS_SDK_TYPES);
             writer.addDependency(TypeScriptDependency.AWS_SDK_PROTOCOL_HTTP);
             // Add the template to each generated test.
             writer.write(IoUtils.readUtf8Resource(getClass(), "protocol-test-stub.ts"));
-            writerInitialized = true;
         }
     }
 

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/TypeScriptDelegator.java

@@ -131,6 +131,16 @@ void useFileWriter(String filename, Consumer<TypeScriptWriter> writerConsumer) {
         writerConsumer.accept(checkoutWriter(filename));
     }
 
+    /**
+     * Gets a previously created writer or creates a new one if needed
+     * and adds a new line if the writer already exists.
+     *
+     * @param filename Name of the file to create.
+     */
+    TypeScriptWriter checkoutFileWriter(String filename) {
+        return checkoutWriter(filename);
+    }
+
     private TypeScriptWriter checkoutWriter(String filename) {
         String formattedFilename = Paths.get(filename).normalize().toString();
         boolean needsNewline = writers.containsKey(formattedFilename);

smithy-typescript-codegen/src/main/java/software/amazon/smithy/typescript/codegen/integration/ProtocolGenerator.java

@@ -17,6 +17,7 @@
 
 import java.util.Collection;
 import java.util.List;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import software.amazon.smithy.codegen.core.CodegenException;
 import software.amazon.smithy.codegen.core.Symbol;
@@ -270,6 +271,7 @@ class GenerationContext {
         private ServiceShape service;
         private SymbolProvider symbolProvider;
         private TypeScriptWriter writer;
+        private Supplier<TypeScriptWriter> writerSupplier;
         private List<TypeScriptIntegration> integrations;
         private String protocolName;
 
@@ -306,11 +308,24 @@ public void setSymbolProvider(SymbolProvider symbolProvider) {
         }
 
         public TypeScriptWriter getWriter() {
+            if (writerSupplier != null && writer == null) {
+                writer = writerSupplier.get();
+            }
             return writer;
         }
 
         public void setWriter(TypeScriptWriter writer) {
             this.writer = writer;
+            if (writer != null) {
+                this.writerSupplier = null;
+            }
+        }
+
+        public void setDeferredWriter(Supplier<TypeScriptWriter> writerSupplier) {
+            this.writerSupplier = writerSupplier;
+            if (writerSupplier != null) {
+                this.writer = null;
+            }
         }
 
         public List<TypeScriptIntegration> getIntegrations() {
@@ -336,6 +351,7 @@ public GenerationContext copy() {
             copy.setService(service);
             copy.setSymbolProvider(symbolProvider);
             copy.setWriter(writer);
+            copy.setDeferredWriter(writerSupplier);
             copy.setIntegrations(integrations);
             copy.setProtocolName(protocolName);
             return copy;