sdks/python/test-suites/direct/common.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.
 */

def pythonVersionSuffix = project.ext.pythonVersion.replace('.', '')
def pythonContainerVersion = project.ext.pythonVersion == '2.7' ? '2' : project.ext.pythonVersion
def runScriptsDir = "${rootDir}/sdks/python/scripts"
// Basic test options for ITs running on Jenkins.
def basicTestOpts = [
    "--capture=no",  // print stdout instantly
    "--numprocesses=8",  // run tests in parallel
    "--timeout=4500", // timeout of whole command execution
    "--color=yes", // console color
    "--log-cli-level=INFO" //log level info
]

task postCommitIT {
  dependsOn 'installGcpTest'

  // Run IT tests with TestDirectRunner in batch in Python 3.
  doLast {
    def batchTests = [
        "apache_beam/examples/wordcount_it_test.py::WordCountIT::test_wordcount_it",
        "apache_beam/io/gcp/pubsub_integration_test.py::PubSubIntegrationTest",
        "apache_beam/io/gcp/big_query_query_to_table_it_test.py::BigQueryQueryToTableIT",
        "apache_beam/io/gcp/bigquery_io_read_it_test.py",
        "apache_beam/io/gcp/bigquery_read_it_test.py",
        "apache_beam/io/gcp/bigquery_write_it_test.py",
        "apache_beam/io/gcp/datastore/v1new/datastore_write_it_test.py",
    ]
    def testOpts = basicTestOpts + ["${batchTests.join(' ')}"]
    def argMap = ["runner": "TestDirectRunner",
                  "test_opts": testOpts,
                  "suite": "postCommitIT-direct-py${pythonVersionSuffix}",
                  ]
    def batchCmdArgs = mapToArgString(argMap)
    exec {
      executable 'sh'
      args '-c', ". ${envdir}/bin/activate && ${runScriptsDir}/run_integration_test.sh $batchCmdArgs"
    }
  }
}

task spannerioIT {
  dependsOn 'installGcpTest'

  // Run Spanner IO IT tests with TestDirectRunner in batch in Python 3.
  doLast {
    def batchTests = [
        "apache_beam/io/gcp/experimental/spannerio_read_it_test.py",
        "apache_beam/io/gcp/experimental/spannerio_write_it_test.py",
    ]
    def testOpts = basicTestOpts + ["${batchTests.join(' ')}"]
    def argMap = ["runner": "TestDirectRunner",
                  "test_opts": testOpts,
                  "suite": "postCommitIT-direct-py${pythonVersionSuffix}",
                  ]
    def batchCmdArgs = mapToArgString(argMap)
    exec {
      executable 'sh'
      args '-c', ". ${envdir}/bin/activate && ${runScriptsDir}/run_integration_test.sh $batchCmdArgs"
    }
  }
}

task examples {
  dependsOn 'installGcpTest'
  dependsOn ':sdks:python:sdist'

  doLast {
    def testOpts = ["--log-cli-level=INFO"]
    def argMap = ["runner": "TestDirectRunner",
                  "test_opts": testOpts,
                  "suite": "postCommitExamples-direct-py${pythonVersionSuffix}",
                  "collect": "examples_postcommit and not sickbay_direct"
    ]
    def batchCmdArgs = mapToArgString(argMap)
    exec {
      executable 'sh'
      args '-c', ". ${envdir}/bin/activate && ${runScriptsDir}/run_integration_test.sh $batchCmdArgs"
    }
  }
}


tasks.register("mongodbioIT") {
  dependsOn 'setupVirtualenv'

  Random r = new Random()
  def port = r.nextInt(1000) + 27017
  def containerName = "mongoioit" + port

  def options = [
          "--mongo_uri=mongodb://localhost:" + port
  ]

  // Pull the latest mongodb docker image and run
  doFirst {
    exec {
      executable 'sh'
      args '-c', "docker pull mongo && docker run --name ${containerName} -p ${port}:27017 -d mongo:latest"
    }
  }

  doLast {
    exec {
      executable 'sh'
      args '-c', ". ${envdir}/bin/activate && pip install -e ${rootDir}/sdks/python/[test] && python -m apache_beam.io.mongodbio_it_test ${options.join(' ')}"
    }
    exec {
      executable 'sh'
      args '-c', "docker stop ${containerName} && docker rm ${containerName}"
    }
  }
}

task directRunnerIT {
  dependsOn 'installGcpTest'

  // Run IT tests with TestDirectRunner in batch.
  doLast {
    def tests = [
        "apache_beam/examples/wordcount_it_test.py::WordCountIT::test_wordcount_it",
        "apache_beam/io/gcp/pubsub_integration_test.py::PubSubIntegrationTest",
        "apache_beam/io/gcp/big_query_query_to_table_it_test.py::BigQueryQueryToTableIT",
        "apache_beam/io/gcp/bigquery_io_read_it_test.py",
        "apache_beam/io/gcp/bigquery_read_it_test.py",
        "apache_beam/io/gcp/bigquery_write_it_test.py",
        "apache_beam/io/gcp/datastore/v1new/datastore_write_it_test.py",
    ]
    def batchTestOpts = basicTestOpts + ["${tests.join(' ')}"]
    def argMap = ["runner": "TestDirectRunner",
                  "test_opts": batchTestOpts,
                  "suite": "directRunnerIT-batch",
                  ]
    def batchCmdArgs = mapToArgString(argMap)
    exec {
      executable 'sh'
      args '-c', ". ${envdir}/bin/activate && ${runScriptsDir}/run_integration_test.sh $batchCmdArgs"
    }
  }

  // Run IT tests with TestDirectRunner in streaming.
  doLast {
    def tests = [
        "apache_beam/examples/wordcount_it_test.py::WordCountIT::test_wordcount_it",
        "apache_beam/io/gcp/pubsub_integration_test.py::PubSubIntegrationTest",
        "apache_beam/io/gcp/bigquery_test.py::BigQueryStreamingInsertTransformIntegrationTests::test_multiple_destinations_transform",
        "apache_beam/io/gcp/bigquery_test.py::PubSubBigQueryIT",
        "apache_beam/io/gcp/bigquery_file_loads_test.py::BigQueryFileLoadsIT::test_bqfl_streaming",
    ]
    def streamingTestOpts = basicTestOpts + ["${tests.join(' ')}"]
    def argMap = ["runner": "TestDirectRunner",
                  "streaming": "true",
                  "test_opts": streamingTestOpts,
                  "suite": "directRunnerIT-streaming",
                ]
    def streamingCmdArgs = mapToArgString(argMap)
    exec {
      executable 'sh'
      args '-c', ". ${envdir}/bin/activate && ${runScriptsDir}/run_integration_test.sh $streamingCmdArgs"
    }
  }
}

task hdfsIntegrationTest {
  doLast {
    exec {
      executable 'sh'
      args '-c', "${rootDir}/sdks/python/apache_beam/io/hdfs_integration_test/hdfs_integration_test.sh python:${pythonContainerVersion}"
    }
  }
}