sdks/python/apache_beam/testing/load_tests/build.gradle

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * License); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an AS IS BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */


apply plugin: org.apache.beam.gradle.BeamModulePlugin
applyPythonNature()

dependencies {
    distTarBall project(path: ":sdks:python", configuration: "distTarBall")
}

description = "Apache Beam :: SDKs :: Python :: Load Tests"

def mainClassProperty = "loadTest.mainClass"
def mainClass = project.findProperty(mainClassProperty)

def loadTestArgsProperty = "loadTest.args"

def runnerProperty = "runner"
def withDataflowWorkerJarProperty = "withDataflowWorkerJar"

task run(type: Exec, dependsOn: installGcpTest) {
    def loadTestArgs = project.findProperty(loadTestArgsProperty) ?: ""

    def runnerArg = project.findProperty(runnerProperty) ?: ""
    if (runnerArg == 'DataflowRunner' || runnerArg == 'TestDataflowRunner') {
        dependsOn ':sdks:python:sdist'
        loadTestArgs +=" --sdk_location=${files(configurations.distTarBall.files).singleFile}"
    }

    String withDataflowWorkerJarArg = project.findProperty(withDataflowWorkerJarProperty).toString() ?: ""
    if (withDataflowWorkerJarArg.toLowerCase() == 'true') {
        dependsOn ':runners:google-cloud-dataflow-java:worker:shadowJar'
        loadTestArgs +=" --dataflow_worker_jar=${project(":runners:google-cloud-dataflow-java:worker").shadowJar.archivePath}"
    }

    setWorkingDir "${project.rootDir}/sdks/python"
    commandLine 'sh', '-c', "${project.ext.envdir}/bin/python -m ${mainClass} --test-pipeline-options=\"${parseOptions(loadTestArgs)}\" --ignore-files \'.*py3\\d?\\.py\$\'"
    ignoreExitValue true

    doLast {
        if (execResult.exitValue != 0) {
            throw new GradleException('error occurred')
        }
    }
}

def parseOptions(String option) {
    option.replace('\"', '\\"')
}