| # Copyright 2022 Google LLC |
| # |
| # 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. |
| |
| """Module for SuiteSessionManager class.""" |
| import datetime |
| import os |
| import time |
| import threading |
| from typing import Any, Callable, Dict, List, Optional |
| |
| from local_agent import errors as agent_errors |
| from local_agent import logger as logger_module |
| |
| |
| logger = logger_module.get_logger() |
| |
| # ============= Constants ============= # |
| _ARTIFACT_EXPIRE_DAYS = 30 # days |
| _SUITE_SESSION_TIME_OUT = 5400 # 90 mins in seconds |
| _SUITE_SESSION_TIME_OUT_HUMAN_READABLE = '90 mins' |
| _SUITE_SESSION_TIME_OUT_INTERVAL_SECONDS = 30 |
| # ===================================== # |
| |
| |
| class SuiteSessionManager: |
| """Suite Session Manager for test suite session management.""" |
| |
| def __init__(self, |
| artifacts_fn: Callable[[str, str], None], |
| artifact_root_dir: str, |
| create_devices_fn: Callable[[List[str], str], None], |
| close_devices_fn: Callable[[], None]): |
| # Artifacts related |
| self._compress_artifacts_and_upload = artifacts_fn |
| self._artifact_root_dir = artifact_root_dir |
| |
| # Tracks test suite related |
| self._termination_event = None |
| self._suite_timeout_checker = threading.Thread( |
| target=self._check_suite_timeout, daemon=True) |
| self._ongoing_test_suite_id = None |
| self._ongoing_test_suite_start_time = None |
| |
| # Device control related |
| self._create_devices = create_devices_fn |
| self._close_devices = close_devices_fn |
| |
| def start(self, termination_event: threading.Event) -> None: |
| """Starts the suite session manager by enabling the background threads. |
| |
| Args: |
| termination_event: The termination threading event for the thread. |
| """ |
| self._termination_event = termination_event |
| self._suite_timeout_checker.start() |
| |
| def start_test_suite(self, rpc_request: Dict[str, str]) -> Dict[str, Any]: |
| """Subroutine for handling START_TEST_SUITE command. |
| |
| Args: |
| rpc_request: JSON-RPC request. |
| |
| Raises: |
| InvalidRPCError: Invalid RPC. |
| InvalidTestSuiteSessionError: Previous test suite has not ended. |
| |
| Returns: |
| RPC response. |
| """ |
| dut_ids = rpc_request['params'].get('dutDeviceIds') |
| if dut_ids is None: |
| raise agent_errors.InvalidRPCError( |
| 'Invalid rpc command, no dutDeviceIds') |
| suite_id = rpc_request['params'].get('id') |
| force_start = rpc_request['params'].get('forceStart') |
| |
| # Support single suite execution only at a time. |
| if self._ongoing_test_suite_id is not None: |
| if force_start: |
| self.clean_up() |
| else: |
| raise agent_errors.InvalidTestSuiteSessionError( |
| f'The previous test suite {self._ongoing_test_suite_id}' |
| ' has not ended yet.') |
| |
| self._ongoing_test_suite_id = suite_id |
| self._ongoing_test_suite_start_time = time.time() |
| test_suite_dir = self._initialize_artifact_directory( |
| self._ongoing_test_suite_id) |
| |
| self._create_devices(dut_ids, test_suite_dir) |
| |
| logger.info(f'Start a new test suite: {self._ongoing_test_suite_id}') |
| return {'id': rpc_request['id'], 'jsonrpc': '2.0', 'result': {}} |
| |
| def end_test_suite(self, rpc_request: Dict[str, str]) -> Dict[str, Any]: |
| """Subroutine for handling END_TEST_SUITE request. |
| |
| Args: |
| rpc_request: JSON-RPC request. |
| |
| Raises: |
| InvalidTestSuiteSessionError: Invalid testsuite name. |
| |
| Returns: |
| RPC response. |
| """ |
| suite_id = rpc_request['params'].get('id') |
| test_result_id = rpc_request['params'].get('testResultId') |
| |
| if suite_id == self._ongoing_test_suite_id: |
| self.clean_up(test_result_id) |
| else: |
| # Invalid testsuite session |
| raise agent_errors.InvalidTestSuiteSessionError( |
| f'Session {suite_id} has never started before.') |
| |
| return {'id': rpc_request['id'], 'jsonrpc': '2.0', 'result': {}} |
| |
| def clean_up(self, test_result_id: Optional[str] = None) -> None: |
| """Cleans up the ongoing test suite and removes outdated artifacts. |
| |
| Args: |
| test_result_id: Test result ID. |
| """ |
| if self._ongoing_test_suite_id is not None: |
| logger.info( |
| f'Cleaning up the test suite: {self._ongoing_test_suite_id}') |
| self._close_devices() |
| self._compress_artifacts_and_upload( |
| test_suite_id=self._ongoing_test_suite_id, |
| test_result_id=test_result_id) |
| |
| elapsed_time = time.time() - self._ongoing_test_suite_start_time |
| logger.info( |
| f'Cleaned up test suite {self._ongoing_test_suite_id}.' |
| f'Suite session elapsed time = {elapsed_time} seconds.') |
| self._ongoing_test_suite_id = None |
| self._ongoing_test_suite_start_time = None |
| |
| self._remove_outdated_artifacts() |
| |
| def _initialize_artifact_directory(self, ongoing_test_suite_id: str) -> str: |
| """Creates a directory for artifacts of a test suite session. |
| |
| Args: |
| ongoing_test_suite_id: ID of the ongoing test suite. |
| |
| Returns: |
| Artifact directory path. |
| """ |
| test_suite_dir = os.path.join( |
| self._artifact_root_dir, ongoing_test_suite_id) |
| if not os.path.exists(test_suite_dir): |
| os.makedirs(test_suite_dir) |
| |
| local_agent_log = os.path.join(test_suite_dir, 'local_agent.log') |
| logger_module.add_file_handler(local_agent_log) |
| |
| logger.info(f'Start collecting log for {ongoing_test_suite_id}') |
| |
| return test_suite_dir |
| |
| def _remove_outdated_artifacts(self) -> None: |
| """Deletes the outdated local artifacts (older than 1 month).""" |
| if not os.path.exists(self._artifact_root_dir): |
| return |
| today = datetime.datetime.today() |
| for name in os.listdir(self._artifact_root_dir): |
| artifact = os.path.join(self._artifact_root_dir, name) |
| if os.path.isfile(artifact): |
| created_time = os.stat(artifact).st_ctime |
| delta = today - datetime.datetime.fromtimestamp(created_time) |
| if delta.days >= _ARTIFACT_EXPIRE_DAYS: |
| try: |
| os.remove(artifact) |
| logger.info( |
| f'Deleted outdated artifact: {artifact}') |
| except OSError as e: |
| logger.warning( |
| 'Failed to remove the outdated artifact' |
| f' {artifact}: {str(e)}') |
| |
| def _check_suite_timeout(self) -> None: |
| """Checks if the current test suite has timed out. |
| |
| Practically no suite in Rainier will be longer than 1 hour, |
| if the current suite has been running over 90 mins, we'll |
| force clean up the current suite session. |
| """ |
| while (self._termination_event is not None and |
| not self._termination_event.wait( |
| _SUITE_SESSION_TIME_OUT_INTERVAL_SECONDS)): |
| |
| if self._ongoing_test_suite_id is None: |
| logger.info('No ongoing test suite.') |
| continue |
| |
| logger.info('Checking if suite session has timed out.') |
| suite_elapsed_time = ( |
| time.time() - self._ongoing_test_suite_start_time) |
| if suite_elapsed_time >= _SUITE_SESSION_TIME_OUT: |
| logger.warning( |
| f'Suite {self._ongoing_test_suite_id} has timed out ' |
| f'(over {_SUITE_SESSION_TIME_OUT_HUMAN_READABLE}), ' |
| 'forcing suite session clean up.') |
| self.clean_up() |
