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 sdkLocationOverrideProperty = "sdkLocationOverride"

def requirementsTxtFileProperty = "loadTest.requirementsTxtFile"

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

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

    String requirementsTxtFileArg = project.findProperty(requirementsTxtFileProperty) ?: null
    if (requirementsTxtFileArg != null) {
    doFirst {
        exec {
          setWorkingDir "${project.rootDir}/sdks/python"
          executable 'sh'
          args '-c', "${project.ext.envdir}/bin/python -m pip install -r ${requirementsTxtFileArg}"
        }
     }
   }
    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 (executionResult.get().exitValue != 0) {
            throw new GradleException('error occurred')
        }
    }
}

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