text: # Copyright 2019-2021 The Kubeflow Authors
  #
  # Licensed 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 kfp
  import os
  import tarfile
  import time
  import utils
  import yamale
  import yaml
  from datetime import datetime
  from kfp import Client
  from constants import CONFIG_DIR, DEFAULT_CONFIG, SCHEMA_CONFIG
  
  
  class PySampleChecker(object):
  
      def __init__(
              self,
              testname,
              input,
              output,
              result,
              experiment_name,
              host,
              namespace='kubeflow'
      ):
          """Util class for checking python sample test running results.
  
      :param testname: test name.
      :param input: The path of a pipeline file that will be submitted.
      :param output: The path of the test output.
      :param result: The path of the test result that will be exported.
      :param host: The hostname of KFP API endpoint.
      :param namespace: namespace of the deployed pipeline system. Default: kubeflow
      :param experiment_name: Name of the experiment to monitor
      """
          self._testname = testname
          self._experiment_name = experiment_name
          self._input = input
          self._output = output
          self._result = result
          self._host = host
          self._namespace = namespace
          self._run_pipeline = None
          self._test_timeout = None
  
          self._test_cases = []
          self._test_name = self._testname + ' Sample Test'
  
          self._client = None
          self._experiment_id = None
          self._job_name = None
          self._test_args = None
          self._run_id = None
  
      def run(self):
          """Run compiled KFP pipeline."""
  
          ###### Initialization ######
          self._client = Client(host=self._host)
  
          ###### Check Input File ######
          utils.add_junit_test(
              self._test_cases, 'input generated yaml file',
              os.path.exists(self._input), 'yaml file is not generated'
          )
          if not os.path.exists(self._input):
              utils.write_junit_xml(
                  self._test_name, self._result, self._test_cases
              )
              print('Error: job not found.')
              exit(1)
  
          ###### Create Experiment ######
          response = self._client.create_experiment(self._experiment_name)
          self._experiment_id = response.id
          utils.add_junit_test(self._test_cases, 'create experiment', True)
  
          ###### Create Job ######
          self._job_name = self._testname + '_sample'
          ###### Figure out arguments from associated config files. #######
          self._test_args = {}
          config_schema = yamale.make_schema(SCHEMA_CONFIG)
          try:
              with open(DEFAULT_CONFIG, 'r') as f:
                  raw_args = yaml.safe_load(f)
              default_config = yamale.make_data(DEFAULT_CONFIG)
              yamale.validate(
                  config_schema, default_config
              )  # If fails, a ValueError will be raised.
          except yaml.YAMLError as yamlerr:
              raise RuntimeError('Illegal default config:{}'.format(yamlerr))
          except OSError as ose:
              raise FileExistsError('Default config not found:{}'.format(ose))
          else:
              self._test_timeout = raw_args['test_timeout']
              self._run_pipeline = raw_args['run_pipeline']
  
          try:
              config_file = os.path.join(
                  CONFIG_DIR, '%s.config.yaml' % self._testname
              )
              with open(config_file, 'r') as f:
                  raw_args = yaml.safe_load(f)
              test_config = yamale.make_data(config_file)
              yamale.validate(
                  config_schema, test_config
              )  # If fails, a ValueError will be raised.
          except yaml.YAMLError as yamlerr:
              print(
                  'No legit yaml config file found, use default args:{}'.
                  format(yamlerr)
              )
          except OSError as ose:
              print(
                  'Config file with the same name not found, use default args:{}'.
                  format(ose)
              )
          else:
              if 'arguments' in raw_args.keys() and raw_args['arguments']:
                  self._test_args.update(raw_args['arguments'])
              if 'output' in self._test_args.keys(
              ):  # output is a special param that has to be specified dynamically.
                  self._test_args['output'] = self._output
              if 'test_timeout' in raw_args.keys():
                  self._test_timeout = raw_args['test_timeout']
              if 'run_pipeline' in raw_args.keys():
                  self._run_pipeline = raw_args['run_pipeline']
  
          # TODO(numerology): Special treatment for TFX::OSS sample
          if self._testname == 'parameterized_tfx_oss':
              self._test_args['pipeline-root'] = os.path.join(
                  self._test_args['output'],
                  'tfx_taxi_simple_' + kfp.dsl.RUN_ID_PLACEHOLDER
              )
              del self._test_args['output']
  
          # Submit for pipeline running.
          if self._run_pipeline:
              response = self._client.run_pipeline(
                  self._experiment_id, self._job_name, self._input,
                  self._test_args
              )
              self._run_id = response.id
              utils.add_junit_test(self._test_cases, 'create pipeline run', True)
  
      def check(self):
          """Check pipeline run results."""
          if self._run_pipeline:
              ###### Monitor Job ######
              try:
                  start_time = datetime.now()
                  response = self._client.wait_for_run_completion(
                      self._run_id, self._test_timeout
                  )
                  succ = (response.run.status.lower() == 'succeeded')
                  end_time = datetime.now()
                  elapsed_time = (end_time - start_time).seconds
                  utils.add_junit_test(
                      self._test_cases, 'job completion', succ,
                      'waiting for job completion failure', elapsed_time
                  )
              finally:
                  ###### Output Argo Log for Debugging ######
                  workflow_json = self._client._get_workflow_json(self._run_id)
                  workflow_id = workflow_json['metadata']['name']
                  print("Argo Workflow Name: ", workflow_id)
                  argo_log, _ = utils.run_bash_command(
                      'argo logs {} -n {}'.format(
                          workflow_id, self._namespace
                      )
                  )
                  print('=========Argo Workflow Log=========')
                  print(argo_log)
  
              if not succ:
                  utils.write_junit_xml(
                      self._test_name, self._result, self._test_cases
                  )
                  exit(1)
  
              ###### Validate the results for specific test cases ######
              if self._testname == 'xgboost_training_cm':
                  # For xgboost sample, check its confusion matrix.
                  cm_tar_path = './confusion_matrix.tar.gz'
                  utils.get_artifact_in_minio(
                      workflow_json, 'confusion-matrix', cm_tar_path,
                      'mlpipeline-ui-metadata'
                  )
                  with tarfile.open(cm_tar_path) as tar_handle:
                      file_handles = tar_handle.getmembers()
                      assert len(file_handles) == 1
  
                      with tar_handle.extractfile(file_handles[0]) as f:
                          cm_data = f.read()
                          utils.add_junit_test(
                              self._test_cases, 'confusion matrix format',
                              (len(cm_data) > 0),
                              'the confusion matrix file is empty'
                          )
  
          ###### Delete Job ######
          #TODO: add deletion when the backend API offers the interface.
  
          ###### Write out the test result in junit xml ######
          utils.write_junit_xml(self._test_name, self._result, self._test_cases)
  