fix(rosetta): infuse incorrectly handles compressed assemblies (#3669)

The behavior of `rosetta infuse` was incorrectly handled before. `infuse` would always overwrite the `.jsii` file with the uncompressed assembly. This PR fixes that behavior by detecting whether or not there is a compressed file in the directory, and compressing if that is the case.

There are two alternative solutions I considered, primarily because I'm concerned that looking up the location of the compressed assembly can be considered a leaky abstraction:
- loading and looking into the contents of `.jsii` again to determine whether or not it is a file redirect. I did not choose this option as it involves additional loading, which will slow things down.
- We can expand the `LoadedAssembly` type to include information on whether or not the assembly was originally compressed. Then pass that into the `replaceAssembly()` function. I ultimately decided against this because it would involve changing the function signature of all `loadAssemblyFromXxx` functions. It's both a breaking change, and unnecessary clutter for a single use case.

Based on the reasoning above, I think what is included in this PR makes the most sense: expose an independent function, `compressedAssemblyExists`, that returns whether or not there is a file located at `SPEC_FILE_NAME_COMPRESSED`. 

---

By submitting this pull request, I confirm that my contribution is made under the terms of the [Apache 2.0 license].

[Apache 2.0 license]: https://www.apache.org/licenses/LICENSE-2.0

Author: kaizencc

packages/@jsii/spec/src/assembly-utils.ts

@@ -14,6 +14,13 @@ import {
 } from './redirect';
 import { validateAssembly } from './validate-assembly';
 
+/**
+ * Returns true if the SPEC_FILE_NAME_COMPRESSED file exists in the directory.
+ */
+export function compressedAssemblyExists(directory: string): boolean {
+  return fs.existsSync(path.join(directory, SPEC_FILE_NAME_COMPRESSED));
+}
+
 /**
  * Finds the path to the SPEC_FILE_NAME file, which will either
  * be the assembly or hold instructions to find the assembly.

packages/jsii-rosetta/lib/jsii/assemblies.ts

@@ -1,5 +1,11 @@
 import * as spec from '@jsii/spec';
-import { loadAssemblyFromFile, loadAssemblyFromPath, findAssemblyFile, writeAssembly } from '@jsii/spec';
+import {
+  compressedAssemblyExists,
+  loadAssemblyFromFile,
+  loadAssemblyFromPath,
+  findAssemblyFile,
+  writeAssembly,
+} from '@jsii/spec';
 import * as crypto from 'crypto';
 import * as fs from 'fs-extra';
 import * as path from 'path';
@@ -228,14 +234,11 @@ export async function allTypeScriptSnippets(
 
 /**
  * Replaces the file where the original assembly file *should* be found with a new assembly file.
+ * Detects whether or not there is a compressed assembly, and if there is, compresses the new assembly also.
  * Recalculates the fingerprint of the assembly to avoid tampering detection.
  */
-export function replaceAssembly(
-  assembly: spec.Assembly,
-  directory: string,
-  { compress = false }: { compress?: boolean } = {},
-) {
-  writeAssembly(directory, _fingerprint(assembly), { compress });
+export function replaceAssembly(assembly: spec.Assembly, directory: string) {
+  writeAssembly(directory, _fingerprint(assembly), { compress: compressedAssemblyExists(directory) });
 }
 
 /**

packages/jsii-rosetta/test/commands/infuse.test.ts

@@ -1,4 +1,4 @@
-import { loadAssemblyFromPath } from '@jsii/spec';
+import { loadAssemblyFromPath, SPEC_FILE_NAME, SPEC_FILE_NAME_COMPRESSED } from '@jsii/spec';
 import * as fs from 'fs-extra';
 import * as path from 'path';
 
@@ -107,3 +107,58 @@ test('can log to output file', async () => {
   expect(stats.isFile()).toBeTruthy();
   expect(stats.size).toBeGreaterThan(0);
 });
+
+test('preserves the assembly compression if present', async () => {
+  // Create an assembly in a temp directory
+  const compAssembly = TestJsiiModule.fromSource(
+    {
+      'index.ts': `
+      export class ClassA {
+        public someMethod() {
+        }
+      }
+      export class ClassB {
+        public argumentMethod(args: BeeArgs) {
+          Array.isArray(args);
+        }
+      }
+      export interface BeeArgs { readonly value: string; readonly nested?: NestedType; }
+      export interface NestedType { readonly x: number; }
+      `,
+      'README.md': DUMMY_README,
+    },
+    {
+      name: 'my_assembly',
+      jsii: DUMMY_JSII_CONFIG,
+    },
+    {
+      compressAssembly: true,
+    },
+  );
+
+  // Ensure assembly is compressed
+  expect(fs.existsSync(path.join(compAssembly.moduleDirectory, SPEC_FILE_NAME_COMPRESSED))).toBeTruthy();
+
+  // Create a tabletFile in the same directory
+  await extractSnippets([compAssembly.moduleDirectory], {
+    cacheToFile: path.join(compAssembly.moduleDirectory, TABLET_FILE),
+    includeCompilerDiagnostics: false,
+    validateAssemblies: false,
+  });
+
+  // Now infuse
+  await infuse([compAssembly.moduleDirectory]);
+
+  // Expect file at SPEC_FILE_NAME to still be a file redirect (not the actual assembly)
+  expect(fs.readJSONSync(path.join(compAssembly.moduleDirectory, SPEC_FILE_NAME))).toEqual({
+    schema: 'jsii/file-redirect',
+    compression: 'gzip',
+    filename: SPEC_FILE_NAME_COMPRESSED,
+  });
+
+  // Infuse works as expected
+  const assemblies = loadAssemblies([compAssembly.moduleDirectory], false);
+  const types = assemblies[0].assembly.types;
+  expect(types).toBeDefined();
+  expect(types!['my_assembly.ClassA'].docs?.example).toBeDefined();
+});