feat(amazonq): Extract ZIP File and Unit Test Cases (#5416)

* Adding unit test cases * Added extractZipFile functionality
aws · Feb 27, 2025 · bd4dd63 · bd4dd63
1 parent 35b0424
commit bd4dd63
Show file tree

Hide file tree

Showing 8 changed files with 443 additions and 61 deletions.
diff --git a/...nity/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/artifacts/ArtifactHelper.kt b/...nity/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/artifacts/ArtifactHelper.kt
@@ -5,6 +5,7 @@ package software.aws.toolkits.jetbrains.services.amazonq.lsp.artifacts
 
 import com.intellij.util.io.createDirectories
 import com.intellij.util.text.SemVer
+import org.jetbrains.annotations.VisibleForTesting
 import software.aws.toolkits.core.utils.deleteIfExists
 import software.aws.toolkits.core.utils.error
 import software.aws.toolkits.core.utils.exists
@@ -45,9 +46,23 @@ class ArtifactHelper(private val lspArtifactsPath: Path = DEFAULT_ARTIFACT_PATH,
     }
 
     fun deleteOlderLspArtifacts(manifestVersionRanges: ArtifactManager.SupportedManifestVersionRange) {
+        val validVersions = getAllLocalLspArtifactsWithinManifestRange(manifestVersionRanges)
+
+        // Keep the latest 2 versions, delete others
+        validVersions.drop(2).forEach { (folder, _) ->
+            try {
+                folder.toFile().deleteRecursively()
+                logger.info { "Deleted older LSP artifact: ${folder.fileName}" }
+            } catch (e: Exception) {
+                logger.error(e) { "Failed to delete older LSP artifact: ${folder.fileName}" }
+            }
+        }
+    }
+
+    fun getAllLocalLspArtifactsWithinManifestRange(manifestVersionRanges: ArtifactManager.SupportedManifestVersionRange): List<Pair<Path, SemVer>> {
         val localFolders = getSubFolders(lspArtifactsPath)
 
-        val validVersions = localFolders
+        return localFolders
             .mapNotNull { localFolder ->
                 SemVer.parseFromText(localFolder.fileName.toString())?.let { semVer ->
                     if (semVer in manifestVersionRanges.startVersion..manifestVersionRanges.endVersion) {
@@ -58,16 +73,6 @@ class ArtifactHelper(private val lspArtifactsPath: Path = DEFAULT_ARTIFACT_PATH,
                 }
             }
             .sortedByDescending { (_, semVer) -> semVer }
-
-        // Keep the latest 2 versions, delete others
-        validVersions.drop(2).forEach { (folder, _) ->
-            try {
-                folder.toFile().deleteRecursively()
-                logger.info { "Deleted older LSP artifact: ${folder.fileName}" }
-            } catch (e: Exception) {
-                logger.error(e) { "Failed to delete older LSP artifact: ${folder.fileName}" }
-            }
-        }
     }
 
     fun getExistingLspArtifacts(versions: List<ManifestManager.Version>, target: ManifestManager.VersionTarget?): Boolean {
@@ -103,23 +108,27 @@ class ArtifactHelper(private val lspArtifactsPath: Path = DEFAULT_ARTIFACT_PATH,
             logger.info { "Attempt ${currentAttempt.get()} of $maxDownloadAttempts to download LSP artifacts" }
 
             try {
-                if (downloadLspArtifacts(temporaryDownloadPath, target)) {
+                if (downloadLspArtifacts(temporaryDownloadPath, target) && target != null && !target.contents.isNullOrEmpty()) {
                     moveFilesFromSourceToDestination(temporaryDownloadPath, downloadPath)
+                    target.contents
+                        .mapNotNull { it.filename }
+                        .forEach { filename -> extractZipFile(downloadPath.resolve(filename), downloadPath) }
                     logger.info { "Successfully downloaded and moved LSP artifacts to $downloadPath" }
                     return
                 }
             } catch (e: Exception) {
                 logger.error(e) { "Failed to download/move LSP artifacts on attempt ${currentAttempt.get()}" }
                 temporaryDownloadPath.toFile().deleteRecursively()
-
-                if (currentAttempt.get() >= maxDownloadAttempts) {
-                    throw LspException("Failed to download LSP artifacts after $maxDownloadAttempts attempts", LspException.ErrorCode.DOWNLOAD_FAILED)
-                }
+                downloadPath.toFile().deleteRecursively()
             }
         }
+        if (currentAttempt.get() >= maxDownloadAttempts) {
+            throw LspException("Failed to download LSP artifacts after $maxDownloadAttempts attempts", LspException.ErrorCode.DOWNLOAD_FAILED)
+        }
     }
 
-    private fun downloadLspArtifacts(downloadPath: Path, target: ManifestManager.VersionTarget?): Boolean {
+    @VisibleForTesting
+    internal fun downloadLspArtifacts(downloadPath: Path, target: ManifestManager.VersionTarget?): Boolean {
         if (target == null || target.contents.isNullOrEmpty()) {
             logger.warn { "No target contents available for download" }
             return false
@@ -166,7 +175,8 @@ class ArtifactHelper(private val lspArtifactsPath: Path = DEFAULT_ARTIFACT_PATH,
         }
     }
 
-    private fun validateFileHash(filePath: Path, expectedHash: String?): Boolean {
+    @VisibleForTesting
+    internal fun validateFileHash(filePath: Path, expectedHash: String?): Boolean {
         if (expectedHash == null) return false
         val contentHash = generateSHA384Hash(filePath)
         return "sha384:$contentHash" == expectedHash

diff --git a/...ity/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/artifacts/ArtifactManager.kt b/...ity/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/artifacts/ArtifactManager.kt
@@ -4,13 +4,17 @@
 package software.aws.toolkits.jetbrains.services.amazonq.lsp.artifacts
 
 import com.intellij.util.text.SemVer
-import org.assertj.core.util.VisibleForTesting
+import org.jetbrains.annotations.VisibleForTesting
 import software.aws.toolkits.core.utils.error
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
 import software.aws.toolkits.jetbrains.services.amazonq.project.manifest.ManifestManager
 
-class ArtifactManager {
+class ArtifactManager(
+    private val manifestFetcher: ManifestFetcher = ManifestFetcher(),
+    private val artifactHelper: ArtifactHelper = ArtifactHelper(),
+    manifestRange: SupportedManifestVersionRange?,
+) {
 
     data class SupportedManifestVersionRange(
         val startVersion: SemVer,
@@ -21,20 +25,7 @@ class ArtifactManager {
         val inRangeVersions: List<ManifestManager.Version>,
     )
 
-    private val manifestFetcher: ManifestFetcher
-    private val artifactHelper: ArtifactHelper
-    private val manifestVersionRanges: SupportedManifestVersionRange
-
-    // Primary constructor with config
-    constructor(
-        manifestFetcher: ManifestFetcher = ManifestFetcher(),
-        artifactFetcher: ArtifactHelper = ArtifactHelper(),
-        manifestRange: SupportedManifestVersionRange?,
-    ) {
-        manifestVersionRanges = manifestRange ?: DEFAULT_VERSION_RANGE
-        this.manifestFetcher = manifestFetcher
-        this.artifactHelper = artifactFetcher
-    }
+    private val manifestVersionRanges: SupportedManifestVersionRange = manifestRange ?: DEFAULT_VERSION_RANGE
 
     // Secondary constructor with no parameters
     constructor() : this(ManifestFetcher(), ArtifactHelper(), null)
@@ -57,7 +48,11 @@ class ArtifactManager {
         this.artifactHelper.removeDelistedVersions(lspVersions.deListedVersions)
 
         if (lspVersions.inRangeVersions.isEmpty()) {
-            // No versions are found which are in the given range.
+            // No versions are found which are in the given range. Fallback to local lsp artifacts.
+            val localLspArtifacts = this.artifactHelper.getAllLocalLspArtifactsWithinManifestRange(manifestVersionRanges)
+            if (localLspArtifacts.isNotEmpty()) {
+                return
+            }
             throw LspException("Language server versions not found in manifest.", LspException.ErrorCode.NO_COMPATIBLE_LSP_VERSION)
         }
 

diff --git a/...munity/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/artifacts/LspException.kt b/...munity/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/artifacts/LspException.kt
@@ -11,6 +11,7 @@ class LspException(message: String, private val errorCode: ErrorCode, cause: Thr
         HASH_MISMATCH,
         TARGET_NOT_FOUND,
         NO_COMPATIBLE_LSP_VERSION,
+        UNZIP_FAILED,
     }
 
     override fun toString(): String = buildString {

diff --git a/...-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/artifacts/LspUtils.kt b/...-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/artifacts/LspUtils.kt
@@ -7,11 +7,16 @@ import com.intellij.openapi.util.SystemInfo
 import com.intellij.openapi.util.text.StringUtil
 import com.intellij.util.io.DigestUtil
 import com.intellij.util.system.CpuArch
+import software.aws.toolkits.core.utils.createParentDirectories
+import software.aws.toolkits.core.utils.exists
+import java.io.FileNotFoundException
+import java.io.FileOutputStream
 import java.nio.file.Files
 import java.nio.file.Path
 import java.nio.file.Paths
 import java.nio.file.StandardCopyOption
 import java.security.MessageDigest
+import java.util.zip.ZipFile
 import kotlin.io.path.isDirectory
 import kotlin.io.path.listDirectoryEntries
 
@@ -66,3 +71,26 @@ fun moveFilesFromSourceToDestination(sourceDir: Path, targetDir: Path) {
         throw IllegalStateException("Failed to move files from $sourceDir to $targetDir", e)
     }
 }
+
+fun extractZipFile(zipFilePath: Path, destDir: Path) {
+    if (!zipFilePath.exists()) {
+        throw FileNotFoundException("Zip file not found: $zipFilePath")
+    }
+
+    try {
+        ZipFile(zipFilePath.toFile()).use { zipFile ->
+            zipFile.entries()
+                .asSequence()
+                .filterNot { it.isDirectory }
+                .map { zipEntry ->
+                    val destPath = destDir.resolve(zipEntry.name)
+                    destPath.createParentDirectories()
+                    FileOutputStream(destPath.toFile()).use { targetFile ->
+                        zipFile.getInputStream(zipEntry).copyTo(targetFile)
+                    }
+                }.toList()
+        }
+    } catch (e: Exception) {
+        throw LspException("Failed to extract zip file: ${e.message}", LspException.ErrorCode.UNZIP_FAILED, cause = e)
+    }
+}
diff --git a/...ity/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/artifacts/ManifestFetcher.kt b/...ity/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/artifacts/ManifestFetcher.kt
@@ -3,7 +3,7 @@
 
 package software.aws.toolkits.jetbrains.services.amazonq.lsp.artifacts
 
-import org.assertj.core.util.VisibleForTesting
+import org.jetbrains.annotations.VisibleForTesting
 import software.aws.toolkits.core.utils.deleteIfExists
 import software.aws.toolkits.core.utils.error
 import software.aws.toolkits.core.utils.exists

diff --git a/...s-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/EncoderServer.kt b/...s-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/EncoderServer.kt
@@ -21,15 +21,13 @@ import com.nimbusds.jwt.JWTClaimsSet
 import com.nimbusds.jwt.SignedJWT
 import org.apache.commons.codec.digest.DigestUtils
 import software.amazon.awssdk.utils.UserHomeDirectoryUtils
-import software.aws.toolkits.core.utils.createParentDirectories
-import software.aws.toolkits.core.utils.exists
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
 import software.aws.toolkits.core.utils.tryDirOp
 import software.aws.toolkits.core.utils.warn
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.artifacts.extractZipFile
 import software.aws.toolkits.jetbrains.services.amazonq.project.manifest.ManifestManager
 import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings
-import java.io.FileOutputStream
 import java.io.IOException
 import java.nio.file.Files
 import java.nio.file.Path
@@ -39,7 +37,6 @@ import java.security.Key
 import java.security.SecureRandom
 import java.util.Base64
 import java.util.concurrent.atomic.AtomicInteger
-import java.util.zip.ZipFile
 import javax.crypto.spec.SecretKeySpec
 
 class EncoderServer(val project: Project) : Disposable {
@@ -183,7 +180,7 @@ class EncoderServer(val project: Project) : Disposable {
             if (serverContent?.url != null) {
                 if (validateHash(serverContent.hashes?.first(), HttpRequests.request(serverContent.url).readBytes(null))) {
                     downloadFromRemote(serverContent.url, zipFilePath)
-                    unzipFile(zipFilePath, cachePath)
+                    extractZipFile(zipFilePath, cachePath)
                 }
             }
         } catch (e: Exception) {
@@ -231,26 +228,6 @@ class EncoderServer(val project: Project) : Disposable {
         Files.setPosixFilePermissions(filePath, permissions)
     }
 
-    private fun unzipFile(zipFilePath: Path, destDir: Path) {
-        if (!zipFilePath.exists()) return
-        try {
-            val zipFile = ZipFile(zipFilePath.toFile())
-            zipFile.use { file ->
-                file.entries().asSequence()
-                    .filterNot { it.isDirectory }
-                    .map { zipEntry ->
-                        val destPath = destDir.resolve(zipEntry.name)
-                        destPath.createParentDirectories()
-                        FileOutputStream(destPath.toFile()).use { targetFile ->
-                            zipFile.getInputStream(zipEntry).copyTo(targetFile)
-                        }
-                    }.toList()
-            }
-        } catch (e: Exception) {
-            logger.warn { "error while unzipping project context artifact: ${e.message}" }
-        }
-    }
-
     private fun downloadFromRemote(url: String, path: Path) {
         try {
             HttpRequests.request(url).saveToFile(path, null)