lib/stage_feature_extractor_integration_test.py

# Copyright 2017 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Intgration tests for bad cl detector application."""

# pylint: disable=g-bad-import-order
from __future__ import print_function

import datetime
import mock
import os
import pandas
import unittest

import constants
import stage_feature_extractor
import feature_extractor_utils
import fetch_data_engine
from protos import StagesForCQ_pb2

from chromite.lib import cros_test_lib
from chromite.lib import osutils


STAGE_FEATURE_DIR = 'data/stage/feature'
STAGE_RAW_DIR = 'data/stage/raw_data'


class StageFeatureExtractorIntegrationTest(cros_test_lib.MockTempDirTestCase):
  """Integration test for stage_feature_extractor.py."""

  def setUp(self):
    """Generate test data."""
    os.chdir(self.tempdir)
    self.PatchObject(feature_extractor_utils, 'GetInOutDir',
                     mock.Mock(return_value=[STAGE_RAW_DIR, STAGE_FEATURE_DIR]))
    # create the folder
    osutils.SafeMakedirs(STAGE_FEATURE_DIR)
    osutils.SafeMakedirs(STAGE_RAW_DIR)

    time_now = datetime.datetime.now()
    # put data in test input folders
    stages_dict_list_1 = [
        {'build_id': 11, 'status': 'fail',
         'finish_time': time_now + datetime.timedelta(seconds=5),
         'start_time': time_now,
         'final': 1, 'important': 0,
         'board': 'test1', 'id': 11, 'build_config': 'config_test1',
         'name': 'Uprev'},
        {'build_id': 12, 'status': 'pass',
         'finish_time': 0, 'start_time': 0,
         'final': 1, 'important': 1,
         'board': 'test', 'id': 12, 'build_config': 'config_test1',
         'name': 'Uprev'},
        {'build_id': 13, 'status': 'fail',
         'finish_time': 0, 'start_time': 0,
         'final': 0, 'important': 1,
         'board': 'test', 'id': 13, 'build_config': 'config_test1',
         'name': 'Uprev'}]

    stages_dict_list_2 = [
        {'build_id': 21, 'status': 'pass',
         'finish_time': 0, 'start_time': 0,
         'final': 1, 'important': 1,
         'board': 'test', 'id': 11, 'build_config': 'config_test2',
         'name': 'Uprev'},
        {'build_id': 22, 'status': 'pass',
         'finish_time': 0, 'start_time': 0,
         'final': 1, 'important': 1,
         'board': 'test', 'id': 12, 'build_config': 'config_test2',
         'name': 'Uprev'},
        {'build_id': 23, 'status': 'fail',
         'finish_time': 0, 'start_time': 0,
         'final': 1, 'important': 1,
         'board': 'test', 'id': 13, 'build_config': 'config_test2',
         'name': 'SignerTest'}]

    self._GenerateStagesForCQMsg(1, 'bad_cl', stages_dict_list_1)
    self._GenerateStagesForCQMsg(2, 'lab_failure', stages_dict_list_2)

  def _GenerateStagesForCQMsg(self, master_build_id, failure_category,
                              stages_dict_list):
    """Generate a StagesForCQ() message.

    Args:
      master_build_id: an int, id of master build in CQ.
      failure_category: string, category of the failure, e.g.: 'bad_cl'.
      stages_dict_list: a list of dictionaries. For each dictionary, available
                        keys are (id, build_id, name, board, status, final,
                        last_updated, start_time, finish_time, build_config,
                        important)
    """
    filename = os.path.join(STAGE_RAW_DIR, '%d_0' % master_build_id)
    stages_for_cq_msg = StagesForCQ_pb2.StagesForCQ()
    stages_for_cq_msg.master_build_id = master_build_id
    stages_for_cq_msg.failure_category = failure_category

    for stage in stages_dict_list:
      fetch_data_engine._FillInStageInfo(stages_for_cq_msg.stage_msgs.add(),
                                         stage)

    fetch_data_engine.WriteRecordToFile(stages_for_cq_msg, filename)

  def testExtractStageFeaturesForTraining(self):
    """Test stage_feature_extractor.ExtractStageFeatures() for training."""
    stage_feature_extractor.ExtractStageFeatures(True)
    # for training, all CQs' stages are in the same STAGE_TRAIN_FILE.
    generated_files = os.listdir(STAGE_FEATURE_DIR)
    self.assertEqual(generated_files, [constants.STAGE_TRAIN_FILE])
    df = pandas.read_csv(
        os.path.join(STAGE_FEATURE_DIR, constants.STAGE_TRAIN_FILE))

    self.assertEqual(df['duration'][0], 5)
    self.assertEqual(df['failure_category'][0], 'bad_cl')
    self.assertEqual(df['Uprev'][0], 1)
    self.assertEqual(df['SignerTest'][2], 1)
    self.assertEqual(df['failure_category'][2], 'lab_failure')

  def testExtractStageFeaturesForPrediction(self):
    """Test stage_feature_extractor.ExtractStageFeatures() for prediction."""
    stage_feature_extractor.ExtractStageFeatures(False)
    # for prediction, each CQ's record is stored in its own csv file.
    generated_files = os.listdir(STAGE_FEATURE_DIR)
    # filename of the form 'masterID_count.csv'
    self.assertEqual(generated_files, ['1_0.csv', '2_0.csv'])

    df_1 = pandas.read_csv(
        os.path.join(STAGE_FEATURE_DIR, '1_0.csv'))

    self.assertEqual(df_1['duration'][0], 5)
    self.assertEqual(df_1['failure_category'][0], 'bad_cl')
    self.assertEqual(df_1['Uprev'][0], 1)

    df_2 = pandas.read_csv(
        os.path.join(STAGE_FEATURE_DIR, '2_0.csv'))

    self.assertEqual(df_2['final'][0], 1)
    self.assertEqual(df_2['SignerTest'][0], 1)
    self.assertEqual(df_2['failure_category'][0], 'lab_failure')


if __name__ == '__main__':
  unittest.main()