teamcity-executors-test-task/src/test/kotlin/backup/BackupClientTest.kt

package backup

import aws.sdk.kotlin.services.s3.*
import kotlinx.coroutines.runBlocking
import org.junit.jupiter.api.*
import org.junit.jupiter.api.Assertions.*
import java.io.File
import kotlin.io.path.*
import kotlin.random.Random

val bucketName = System.getenv("BACKUP_BUCKET") ?: "teamcity-executors-test-task"

class BackupClientTest {
    lateinit var s3: S3Client
    lateinit var backupClient: BackupClient

    @BeforeEach
    fun `before each`() =
        runBlocking {
            s3 = S3Client.fromEnvironment {}
            backupClient = BackupClient(s3, bucketName, 1024 * 1024 * 10)
        }

    @AfterEach
    fun `after each`() {
        s3.close()
    }

    @TestFactory
    fun `round-trip tests`() =
        listOf(
            "empty directory" to {
                listOf(
                    Path("_test").createDirectory(),
                )
            },
            "single file" to {
                listOf(
                    Path("_test.txt").apply { writeText("Hello World!") },
                )
            },
            "directory structure" to {
                listOf(
                    Path("_test").createDirectory(),
                    Path("_test/a.txt").apply { writeText("This is file A!\nAnother line here.") },
                    Path("_test/folder").createDirectory(),
                    Path("_test/another-folder").createDirectory(),
                    Path("_test/another-folder/b").apply { writeText("This is file B\n") },
                    Path("_test/another-folder/c.txt").createFile(),
                    Path("_test/README.md").apply { writeText("# This is a test directory structure.") },
                )
            },
            "single large file" to {
                val bytes = ByteArray(1024 * 1024 * 32)
                Random.nextBytes(bytes)
                listOf(
                    Path("_test.txt").apply { writeBytes(bytes) },
                )
            },
            "large directory structure" to {
                val bytes1 = ByteArray(1024 * 1024 * 32)
                val bytes2 = ByteArray(1024 * 1024 * 48)
                listOf(
                    Path("_test").createDirectory(),
                    Path("_test/a.txt").apply { writeBytes(bytes1) },
                    Path("_test/folder").createDirectory(),
                    Path("_test/another-folder").createDirectory(),
                    Path("_test/another-folder/b").apply { writeText("This is file B\n") },
                    Path("_test/another-folder/c.txt").createFile(),
                    Path("_test/README.md").apply { writeBytes(bytes2) },
                )
            },
        ).map { (name, pathsGen) ->
            DynamicTest.dynamicTest(name) {
                val paths = pathsGen()
                val backupKey =
                    assertDoesNotThrow("should upload files") {
                        runBlocking {
                            backupClient.upload(paths.first().toFile())
                        }
                    }
                val restoreDir = Path("_test_restore").createDirectory()
                assertDoesNotThrow("should recover files") {
                    runBlocking {
                        backupClient.restore(restoreDir, backupKey)
                    }
                }
                assertEquals(
                    paths.size,
                    restoreDir.toFile().countEntries() - 1,
                    "number of files in backup restore should be equal to original",
                )
                val individualRestoreDir = Path("_test_restore_individual").createDirectory()
                paths.forEach { path ->
                    if (path.isDirectory()) {
                        individualRestoreDir.resolve(path).createDirectory()
                    } else {
                        assertDoesNotThrow("should recover file '$path'") {
                            runBlocking {
                                backupClient.restoreFile(
                                    individualRestoreDir.resolve(path).parent,
                                    backupKey,
                                    path.toString(),
                                )
                            }
                        }
                    }
                }
                assertEquals(
                    paths.size,
                    individualRestoreDir.toFile().countEntries() - 1,
                    "number of files in individual backup restore should be equal to original",
                )
                paths.asReversed().forEach { path ->
                    val restorePath = restoreDir.resolve(path)
                    val individualRestorePath = individualRestoreDir.resolve(path)
                    if (path.isDirectory()) {
                        assertTrue(restorePath.exists(), "'$path' should exist in backup")
                        assertTrue(restorePath.isDirectory(), "'$path' should be a directory in backup")
                        assertTrue(individualRestorePath.exists(), "'$path' should exist in backup (individual)")
                        assertTrue(
                            individualRestorePath.isDirectory(),
                            "'$path' should be a directory in backup (individual)",
                        )
                    } else {
                        val originalBytes = path.toFile().readBytes().asList()
                        assertEquals(
                            originalBytes,
                            restorePath.toFile().readBytes().asList(),
                            "File contents of '$path' should equal",
                        )
                        assertEquals(
                            originalBytes,
                            individualRestorePath.toFile().readBytes().asList(),
                            "File contents of '$path' should equal (individual)",
                        )
                    }
                    // cleanup
                    path.deleteExisting()
                    restorePath.deleteExisting()
                    individualRestorePath.deleteExisting()
                }
                // cleanup
                restoreDir.deleteExisting()
                individualRestoreDir.deleteExisting()
                runBlocking {
                    s3.deleteObject {
                        bucket = bucketName
                        key = backupKey
                    }
                }
            }
        }
}

internal fun File.countEntries(): Int {
    val queue = ArrayDeque<File>()
    queue.add(this)
    return queue.count { file ->
        if (file.isDirectory) {
            queue.addAll(file.listFiles()!!) // shouldn't ever be null, since we know it's a directory
        }
        true
    }
}