apply plugin: DockerRemoteApiPlugin

/*
 * Copyright OpenSearch Contributors
 * SPDX-License-Identifier: Apache-2.0
 */

import com.bmuschko.gradle.docker.DockerRemoteApiPlugin
import com.bmuschko.gradle.docker.tasks.container.DockerCreateContainer
import com.bmuschko.gradle.docker.tasks.container.DockerRemoveContainer
import com.bmuschko.gradle.docker.tasks.container.DockerStartContainer
import com.bmuschko.gradle.docker.tasks.container.DockerStopContainer
import com.bmuschko.gradle.docker.tasks.image.DockerBuildImage
import com.bmuschko.gradle.docker.tasks.image.DockerPullImage
import com.bmuschko.gradle.docker.tasks.image.Dockerfile
import com.bmuschko.gradle.docker.tasks.network.DockerCreateNetwork
import com.bmuschko.gradle.docker.tasks.network.DockerRemoveNetwork

/**
 * End-to-end test docker network
 */
task createDataPrepperNetwork(type: DockerCreateNetwork) {
    networkName = "data_prepper_network"
}

task removeDataPrepperNetwork(type: DockerRemoveNetwork) {
    dependsOn createDataPrepperNetwork
    networkId = createDataPrepperNetwork.getNetworkId()
}

def AGGREGATE_PIPELINE_YAML = "aggregate-e2e-pipeline.yml"
def LOG_METRICS_PIPELINE_YAML = "log-metrics-pipeline.yml"
def DATA_PREPPER_CONFIG_LOCAL_NODE = "data_prepper_local_node.yml"
def DATA_PREPPER_CONFIG_STATIC = "data_prepper_static.yml"

/**
 * DataPrepper Docker tasks
 */
task createDataPrepperDockerFile(type: Dockerfile) {
    dependsOn copyDataPrepperJar
    destFile = project.file('build/docker/Dockerfile')
    from(dataPrepperBaseImage)
    exposePort(2021)
    workingDir("/app/data-prepper")
    copyFile("${dataPrepperJarFilepath}", "/app/data-prepper/lib")
    copyFile("src/integrationTest/resources/default_certificate.pem", "/app/data-prepper/config/default_certificate.pem")
    copyFile("src/integrationTest/resources/default_private_key.pem", "/app/data-prepper/config/default_private_key.pem")
    defaultCommand('java', '-Ddata-prepper.dir=/app/data-prepper', '-cp', '/app/data-prepper/lib/*', 'org.opensearch.dataprepper.DataPrepperExecute')
}

task buildDataPrepperDockerImage(type: DockerBuildImage) {
    dependsOn createDataPrepperDockerFile
    inputDir  = file(".")
    dockerFile  = file("build/docker/Dockerfile")
    images.add("integ-test-pipeline-image")
}

def createDataPrepperDockerContainer(final String taskBaseName, final String dataPrepperName, final int sourcePort,
                                     final String pipelineConfigYAML, final String dataPrepperConfigYAML) {
    return tasks.create("create${taskBaseName}", DockerCreateContainer) {
        dependsOn buildDataPrepperDockerImage
        dependsOn createDataPrepperNetwork
        containerName = dataPrepperName
        hostConfig.portBindings = [String.format('%d:2021', sourcePort)]
        hostConfig.network = createDataPrepperNetwork.getNetworkName()
        hostConfig.binds = [(project.file("src/integrationTest/resources/${pipelineConfigYAML}").toString()):"/app/data-prepper/pipelines/pipelines.yaml",
                            "/tmp":"/tmp",
                            (project.file("src/integrationTest/resources/${dataPrepperConfigYAML}").toString()):"/app/data-prepper/config/data-prepper-config.yaml"]
        cmd = ['java', '-Ddata-prepper.dir=/app/data-prepper', '-cp', '/app/data-prepper/lib/*', 'org.opensearch.dataprepper.DataPrepperExecute']
        targetImageId buildDataPrepperDockerImage.getImageId()
    }
}

def startDataPrepperDockerContainer(final DockerCreateContainer createDataPrepperDockerContainerTask) {
    return tasks.create("start${createDataPrepperDockerContainerTask.getName()}", DockerStartContainer) {
        dependsOn createDataPrepperDockerContainerTask
        targetContainerId createDataPrepperDockerContainerTask.getContainerId()
    }
}

def stopDataPrepperDockerContainer(final DockerStartContainer startDataPrepperDockerContainerTask) {
    return tasks.create("stop${startDataPrepperDockerContainerTask.getName()}", DockerStopContainer) {
        targetContainerId startDataPrepperDockerContainerTask.getContainerId()
    }
}

def removeDataPrepperDockerContainer(final DockerStopContainer stopDataPrepperDockerContainerTask) {
    return tasks.create("remove${stopDataPrepperDockerContainerTask.getName()}", DockerRemoveContainer) {
        targetContainerId stopDataPrepperDockerContainerTask.getContainerId()
    }
}

/**
 * OpenSearch Docker tasks
 */
task pullOpenSearchDockerImage(type: DockerPullImage) {
    image = "opensearchproject/opensearch:${libs.versions.opensearch.get()}"
}

task createOpenSearchDockerContainer(type: DockerCreateContainer) {
    dependsOn createDataPrepperNetwork
    dependsOn pullOpenSearchDockerImage
    targetImageId pullOpenSearchDockerImage.image
    containerName = "node-0.example.com"
    hostConfig.portBindings = ['9200:9200', '9600:9600']
    hostConfig.autoRemove = true
    hostConfig.network = createDataPrepperNetwork.getNetworkName()
    envVars = ['discovery.type':'single-node']
}

task startOpenSearchDockerContainer(type: DockerStartContainer) {
    dependsOn createOpenSearchDockerContainer
    targetContainerId createOpenSearchDockerContainer.getContainerId()

    doLast {
        sleep(90*1000)
    }
}

task stopOpenSearchDockerContainer(type: DockerStopContainer) {
    targetContainerId createOpenSearchDockerContainer.getContainerId()

    doLast {
        sleep(5*1000)
    }
}

/**
 * End to end test. Spins up OpenSearch and DataPrepper docker containers, then runs the integ test
 * Stops the docker containers when finished
 */
def createEndToEndTest(final String testName, final String includeTestsMatchPattern,
                       final DockerCreateContainer createDataPrepper1Task, final DockerCreateContainer createDataPrepper2Task) {
    return tasks.create(testName, Test) {
        dependsOn build
        dependsOn startOpenSearchDockerContainer
        def startDataPrepper1Task = startDataPrepperDockerContainer(createDataPrepper1Task as DockerCreateContainer)
        def startDataPrepper2Task = startDataPrepperDockerContainer(createDataPrepper2Task as DockerCreateContainer)
        dependsOn startDataPrepper1Task
        dependsOn startDataPrepper2Task
        startDataPrepper1Task.mustRunAfter 'startOpenSearchDockerContainer'
        startDataPrepper2Task.mustRunAfter 'startOpenSearchDockerContainer'
        // wait for data-preppers to be ready
        doFirst {
            sleep(15*1000)
        }

        description = 'Runs the raw span integration tests.'
        group = 'verification'
        testClassesDirs = sourceSets.integrationTest.output.classesDirs
        classpath = sourceSets.integrationTest.runtimeClasspath

        filter {
            includeTestsMatching includeTestsMatchPattern
        }

        finalizedBy stopOpenSearchDockerContainer
        def stopDataPrepper1Task = stopDataPrepperDockerContainer(startDataPrepper1Task as DockerStartContainer)
        def stopDataPrepper2Task = stopDataPrepperDockerContainer(startDataPrepper2Task as DockerStartContainer)
        finalizedBy stopDataPrepper1Task
        finalizedBy stopDataPrepper2Task
        finalizedBy removeDataPrepperDockerContainer(stopDataPrepper1Task as DockerStopContainer)
        finalizedBy removeDataPrepperDockerContainer(stopDataPrepper2Task as DockerStopContainer)
        finalizedBy removeDataPrepperNetwork
    }
}

// Discovery mode: LOCAL_NODE
def includeLocalAggregateTestsMatchPattern = "org.opensearch.dataprepper.integration.peerforwarder.EndToEndPeerForwarderTest.testAggregatePipelineWithSingleNodeEndToEnd*"

def createLocalAggregateDataPrepper1Task = createDataPrepperDockerContainer(
        "localAggregateDataPrepper1", "dataprepper1", 2021, "${AGGREGATE_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_LOCAL_NODE}")
def createLocalAggregateDataPrepper2Task = createDataPrepperDockerContainer(
        "localAggregateDataPrepper2", "dataprepper2", 2022, "${AGGREGATE_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_LOCAL_NODE}")

def localAggregateEndToEndTest = createEndToEndTest("localAggregateEndToEndTest", includeLocalAggregateTestsMatchPattern,
        createLocalAggregateDataPrepper1Task, createLocalAggregateDataPrepper2Task)

// Discovery mode: STATIC with SSL & mTLS
def includeStaticAggregateTestsMatchPattern = "org.opensearch.dataprepper.integration.peerforwarder.EndToEndPeerForwarderTest.testAggregatePipelineWithMultipleNodesEndToEnd*"

def createAggregateDataPrepper1Task = createDataPrepperDockerContainer(
        "staticAggregateDataPrepper1", "node1.data-prepper.example.com", 2021, "${AGGREGATE_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_STATIC}")
def createAggregateDataPrepper2Task = createDataPrepperDockerContainer(
        "staticAggregateDataPrepper2", "node2.data-prepper.example.com", 2022, "${AGGREGATE_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_STATIC}")

def staticAggregateEndToEndTest = createEndToEndTest("staticAggregateEndToEndTest", includeStaticAggregateTestsMatchPattern,
        createAggregateDataPrepper1Task, createAggregateDataPrepper2Task)

// Discovery mode: STATIC - Log pipeline metrics e2e test
def includeLocalLogMetricsTestsMatchPattern = "org.opensearch.dataprepper.integration.peerforwarder.EndToEndLogMetricsTest.testLogMetricsPipelineWithMultipleNodesEndToEnd*"

def createStaticLogMetricsDataPrepper1Task = createDataPrepperDockerContainer(
        "staticLogMetricsDataPrepper1", "node1.data-prepper.example.com", 2021, "${LOG_METRICS_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_STATIC}")
def createStaticLogMetricsDataPrepper2Task = createDataPrepperDockerContainer(
        "staticLogMetricsDataPrepper2", "node2.data-prepper.example.com", 2022, "${LOG_METRICS_PIPELINE_YAML}", "${DATA_PREPPER_CONFIG_STATIC}")

def staticLogMetricsEndToEndTest = createEndToEndTest("staticLogMetricsEndToEndTest", includeLocalLogMetricsTestsMatchPattern,
        createStaticLogMetricsDataPrepper1Task, createStaticLogMetricsDataPrepper2Task)

dependencies {
    integrationTestImplementation project(':data-prepper-api')
    integrationTestImplementation project(':data-prepper-plugins:common')
    integrationTestImplementation project(':data-prepper-plugins:log-generator-source')
    integrationTestImplementation project(':data-prepper-plugins:opensearch')
    integrationTestImplementation project(':data-prepper-plugins:aws-plugin-api')
    integrationTestImplementation project(':data-prepper-plugins:otel-trace-group-processor')
    integrationTestImplementation 'org.awaitility:awaitility:4.2.0'
    integrationTestImplementation libs.opentelemetry.proto
    integrationTestImplementation libs.protobuf.util
    integrationTestImplementation libs.armeria.core
    integrationTestImplementation libs.armeria.grpc
    integrationTestImplementation libs.opensearch.rhlc
    integrationTestImplementation 'com.fasterxml.jackson.core:jackson-databind'
}