/*
 * 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.
 */

import CommonJobProperties as commonJobProperties
import InfluxDBCredentialsHelper

def now = new Date().format("MMddHHmmss", TimeZone.getTimeZone('UTC'))

// Common pipeline args for Dataflow job.
def dataflowPipelineArgs = [
  project         : 'apache-beam-testing',
  region          : 'us-central1',
  staging_location: 'gs://temp-storage-for-end-to-end-tests/staging-it',
  temp_location   : 'gs://temp-storage-for-end-to-end-tests/temp-it',
]

testConfigurations = []
pythonVersions = ['27', '37']

for (pythonVersion in pythonVersions) {
  def taskVersion = pythonVersion == '27' ? '2' : pythonVersion
  testConfigurations.add([
    jobName           : "beam_PerformanceTests_WordCountIT_Py${pythonVersion}",
    jobDescription    : "Python SDK Performance Test - Run WordCountIT in Py${pythonVersion} with 1Gb files",
    jobTriggerPhrase  : "Run Python${pythonVersion} WordCountIT Performance Test",
    test              : "apache_beam.examples.wordcount_it_test:WordCountIT.test_wordcount_it",
    gradleTaskName    : ":sdks:python:test-suites:dataflow:py${taskVersion}:runPerformanceTest",
    pipelineOptions   : dataflowPipelineArgs + [
      job_name             : "performance-tests-wordcount-python${pythonVersion}-batch-1gb${now}",
      runner               : 'TestDataflowRunner',
      publish_to_big_query : true,
      metrics_dataset      : 'beam_performance',
      metrics_table        : "wordcount_py${pythonVersion}_pkb_results",
      influx_measurement   : "wordcount_py${pythonVersion}_results",
      influx_db_name       : InfluxDBCredentialsHelper.InfluxDBDatabaseName,
      influx_hostname      : InfluxDBCredentialsHelper.InfluxDBHostname,
      input                : "gs://apache-beam-samples/input_small_files/ascii_sort_1MB_input.0000*", // 1Gb
      output               : "gs://temp-storage-for-end-to-end-tests/py-it-cloud/output",
      expect_checksum      : "ea0ca2e5ee4ea5f218790f28d0b9fe7d09d8d710",
      num_workers          : '10',
      autoscaling_algorithm: "NONE",  // Disable autoscale the worker pool.
    ]
  ])
}

for (testConfig in testConfigurations) {
  createPythonPerformanceTestJob(testConfig)
}

private void createPythonPerformanceTestJob(Map testConfig) {
  // This job runs the Beam Python performance tests
  job(testConfig.jobName) {
    // Set default Beam job properties.
    commonJobProperties.setTopLevelMainJobProperties(delegate)

    InfluxDBCredentialsHelper.useCredentials(delegate)

    // Run job in postcommit, don't trigger every push.
    commonJobProperties.setAutoJob(delegate, 'H */6 * * *')

    // Allows triggering this build against pull requests.
    commonJobProperties.enablePhraseTriggeringFromPullRequest(
        delegate,
        testConfig.jobDescription,
        testConfig.jobTriggerPhrase,
        )

    publishers {
      archiveJunit('**/nosetests*.xml')
    }

    steps {
      gradle {
        rootBuildScriptDir(commonJobProperties.checkoutDir)
        switches("--info")
        switches("-Ptest-pipeline-options=\"${commonJobProperties.mapToArgString(testConfig.pipelineOptions)}\"")
        switches("-Ptest=${testConfig.test}")
        tasks(testConfig.gradleTaskName)
        commonJobProperties.setGradleSwitches(delegate)
      }
    }
  }
}