tools/test_selection/decisiongraph_invoker_test.py

#!/usr/bin/env vpython3
# Copyright 2025 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import unittest
import json
import sys
import os
import tempfile
import shutil
from unittest.mock import patch, MagicMock

# We need to adjust the path to import the script directly.
# This assumes the test script is in the same directory as the main script.
sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))

import decisiongraph_invoker as dgi
import requests


class TestDecisionGraphInvoker(unittest.TestCase):
  def setUp(self):
    self.maxDiff = None
    self.test_dir = tempfile.mkdtemp()
    self.config_data = {
      'build_id': '12345',
      'change': 123456,
      'patchset': 1,
      'builder': 'chromium-commit-queue',
      'api_key': 'test_api_key_123',
    }
    self.config_file_path = os.path.join(self.test_dir, 'test_config.json')
    self.api_url_with_key = f"{dgi.API_URL}?key={self.config_data['api_key']}"

  def tearDown(self):
    shutil.rmtree(self.test_dir)

  def _create_mock_config_file(self, data):
    """Helper to create a mock JSON config file."""
    with open(self.config_file_path, 'w', encoding='utf-8') as f:
      json.dump(data, f)

  @patch('builtins.print')
  def test_load_config_from_json_success(self, mock_print):
    self._create_mock_config_file(self.config_data)
    config = dgi.load_config_from_json(self.config_file_path)
    self.assertEqual(config, self.config_data)

  @patch('sys.exit')
  @patch('builtins.print')
  def test_load_config_from_json_file_not_found(self, mock_print, mock_exit):
    dgi.load_config_from_json('non_existent_file.json')
    mock_print.assert_any_call(
      'Error: Configuration file not found at non_existent_file.json'
    )
    mock_exit.assert_called_with(1)

  @patch('sys.exit')
  @patch('builtins.print')
  def test_load_config_from_json_missing_args(self, mock_print, mock_exit):
    incomplete_config = {'build_id': '123', 'change': 1}
    self._create_mock_config_file(incomplete_config)
    dgi.load_config_from_json(self.config_file_path)
    mock_print.assert_any_call(
      'Error: Missing required arguments in JSON config file.'
    )
    mock_exit.assert_called_with(1)

  @patch('sys.exit')
  @patch('builtins.print')
  def test_load_config_from_json_invalid_type(self, mock_print, mock_exit):
    invalid_config = self.config_data.copy()
    invalid_config['change'] = 'not_an_int'
    self._create_mock_config_file(invalid_config)
    dgi.load_config_from_json(self.config_file_path)
    mock_print.assert_any_call(
      "Error: Invalid data type for 'change' or 'patchset' in JSON. "
      "Expected integers."
    )
    mock_exit.assert_called_with(1)

  @patch('builtins.print')
  @patch('requests.post')
  def test_fetch_api_data_success(self, mock_post, mock_print):
    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_response.json.return_value = {'status': 'success'}
    mock_response.text = '{"status": "success"}'
    mock_response.raise_for_status.return_value = None

    mock_post.return_value = mock_response

    response = dgi.fetch_api_data(self.api_url_with_key, {})
    self.assertEqual(response, {'status': 'success'})
    mock_print.assert_any_call('{"status": "success"}')
    mock_print.assert_any_call(200)

  @patch('builtins.print')
  @patch('requests.post')
  def test_fetch_api_data_failure_http_error(self, mock_post, mock_print):
    mock_response = MagicMock()
    mock_response.status_code = 500
    mock_response.text = 'Internal Server Error'
    mock_response.raise_for_status.side_effect = requests.exceptions.HTTPError(
      '500 Server Error'
    )

    mock_post.return_value = mock_response

    response = dgi.fetch_api_data(self.api_url_with_key, {})
    self.assertIsNone(response)
    mock_print.assert_any_call('An error occurred: 500 Server Error')
    mock_print.assert_any_call('Internal Server Error')
    mock_print.assert_any_call(500)

  def test_overwrite_filter_file_success(self):
    test_suite = 'browser_tests'
    tests_to_skip = ['Test1.testA', 'Test2.testB']
    result = dgi.overwrite_filter_file(self.test_dir, test_suite, tests_to_skip)
    self.assertTrue(result)
    filter_file_path = os.path.join(self.test_dir, f'{test_suite}.filter')
    self.assertTrue(os.path.exists(filter_file_path))
    with open(filter_file_path, 'r', encoding='utf-8') as f:
      content = f.read()
      self.assertIn('-Test1.testA', content)
      self.assertIn('-Test2.testB', content)

  def test_overwrite_filter_file_appends_wildcard(self):
    """Tests that a wildcard '*' is appended to each test name."""
    test_suite = 'browser_tests'
    tests_to_skip = ['TestClassOne', 'TestClassTwo']
    result = dgi.overwrite_filter_file(self.test_dir, test_suite, tests_to_skip)

    self.assertTrue(result)
    filter_file_path = os.path.join(self.test_dir, f'{test_suite}.filter')

    with open(filter_file_path, 'r', encoding='utf-8') as f:
      content = f.read()

    expected_content = (
      '# A list of tests to be skipped, generated by test'
      ' selection.\n'
      '-TestClassOne.*\n'
      '-TestClassTwo.*\n'
    )
    self.assertEqual(content, expected_content)

  @patch('builtins.print')
  def test_overwrite_filter_file_dir_not_found(self, mock_print):
    result = dgi.overwrite_filter_file('non_existent_dir', 'suite', [])
    self.assertFalse(result)
    mock_print.assert_called_with(
      'Error: Filter file directory not found at non_existent_dir'
    )

  @patch('sys.exit')
  @patch('decisiongraph_invoker.load_config_from_json')
  @patch('decisiongraph_invoker.fetch_api_data')
  def test_main_trigger_phase_success(
    self, mock_fetch, mock_load_config, mock_exit
  ):
    mock_load_config.return_value = self.config_data
    mock_fetch.return_value = {'output': [{'result': 'success'}]}

    with patch(
      'sys.argv',
      [
        'decisiongraph_invoker.py',
        '--test-targets',
        'browser_tests',
        '--sts-config-file',
        self.config_file_path,
        '--test-selection-phase',
        'TRIGGER',
      ],
    ):
      dgi.main()
      mock_fetch.assert_called_once()
      _, kwargs = mock_fetch.call_args
      payload = kwargs['json_payload']
      self.assertTrue(
        payload['graph']['stages'][0]['execution_options']['prepare']
      )
      mock_exit.assert_called_with(0)

  @patch('sys.exit')
  @patch('decisiongraph_invoker.load_config_from_json')
  @patch('decisiongraph_invoker.fetch_api_data')
  def test_main_fetch_phase_success(
    self, mock_fetch, mock_load_config, mock_exit
  ):
    mock_load_config.return_value = self.config_data
    mock_fetch.return_value = {
      'outputs': [
        {
          'checks': [
            {
              'identifier': {'luciTest': {'testSuite': 'browser_tests'}},
              'children': [
                {
                  'identifier': {
                    'luciTest': {'testId': 'TestClass.testExample1'}
                  }
                },
                {
                  'identifier': {
                    'luciTest': {'testId': 'TestClass.testExample2'}
                  }
                },
              ],
            }
          ]
        }
      ]
    }

    with patch(
      'sys.argv',
      [
        'decisiongraph_invoker.py',
        '--test-targets',
        'browser_tests',
        '--sts-config-file',
        self.config_file_path,
        '--test-selection-phase',
        'FETCH',
        '--filter-file-dir',
        self.test_dir,
      ],
    ):
      dgi.main()

      mock_fetch.assert_called_once()
      _, kwargs = mock_fetch.call_args
      payload = kwargs['json_payload']
      self.assertFalse(
        payload['graph']['stages'][0]['execution_options']['prepare']
      )

      filter_file_path = os.path.join(self.test_dir, 'browser_tests.filter')
      self.assertTrue(os.path.exists(filter_file_path))
      with open(filter_file_path, 'r', encoding='utf-8') as f:
        content = f.read()
        self.assertIn('-TestClass.testExample1', content)
        self.assertIn('-TestClass.testExample2', content)

      mock_exit.assert_called_with(0)

  @patch('sys.exit')
  @patch('decisiongraph_invoker.load_config_from_json')
  @patch('decisiongraph_invoker.fetch_api_data')
  def test_main_fetch_phase_key_error(
    self, mock_fetch, mock_load_config, mock_exit
  ):
    mock_load_config.return_value = self.config_data
    # Malformed response missing 'testId'
    mock_fetch.return_value = {
      'outputs': [
        {
          'checks': [
            {
              'identifier': {'luciTest': {'testSuite': 'browser_tests'}},
              'children': [{'identifier': {'luciTest': {}}}],
            }
          ]
        }
      ]
    }

    with patch(
      'sys.argv',
      [
        'decisiongraph_invoker.py',
        '--test-targets',
        'browser_tests',
        '--sts-config-file',
        self.config_file_path,
        '--test-selection-phase',
        'FETCH',
        '--filter-file-dir',
        self.test_dir,
      ],
    ):
      with self.assertRaises(KeyError):
        dgi.main()

  @patch('argparse.ArgumentParser.error')
  def test_main_fetch_phase_missing_dir_arg(self, mock_parser_error):
    # Make the mock raise an exception to halt execution, which is what
    # parser.error() does internally by calling sys.exit().
    mock_parser_error.side_effect = SystemExit

    with self.assertRaises(SystemExit):
      with patch(
        'sys.argv',
        [
          'decisiongraph_invoker.py',
          '--test-targets',
          'browser_tests',
          '--sts-config-file',
          self.config_file_path,
          '--test-selection-phase',
          'FETCH',
        ],
      ):
        dgi.main()

    # Verify that parser.error was called with the correct message.
    mock_parser_error.assert_called_with(
      '--filter-file-dir is required when phase is FETCH.'
    )


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