android_device.py

#!/usr/bin/env python3
#
# Copyright (C) 2025 The Android Open Source Project
#
# 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.
#

from __future__ import annotations

import time
import sqlite3
from typing import Optional, TYPE_CHECKING

from adb import Adb, AdbError
from runnables.apply_device_configs import runnable_apply_device_configs
from runnables.prompt_device_serial import runnable_prompt_device_serial
from runnables.run_system_probe import runnable_run_system_probe
from runnables.run_workload import runnable_run_workload
from utils_terminal import TerminalPrintColors

if TYPE_CHECKING:
    from run_config import RunConfig


class AndroidDevice:
    """A connected android device and its interactions."""

    def __init__(
        self,
        database_connection: sqlite3.Connection,
        device_name: str = "",
        serial: str = "",
    ) -> None:
        """Initializes the AndroidDevice object."""
        self.device_name = device_name
        self._serial = ""
        self._environ_prepared = False
        self.page_size: int = 0
        self.build_fingerprint = ""
        self._adb: Optional[Adb] = None
        self.serial = serial
        self.database_connection = database_connection
        self._record_device_details_in_sqlite_db()
        self.post_boot_logcat_log: Optional[str] = None

    @property
    def serial(self) -> str:
        """The serial of the android device."""
        return self._serial

    @serial.setter
    def serial(self, serial: str):
        self._serial = serial
        if self._serial:
            self._adb = Adb(self._serial)
            self._set_page_size()
            self._set_build_fingerprint()

    def _set_page_size(self) -> None:
        if not self._adb:
            return
        try:
            result = self._adb.adb_command(["shell", "getconf", "PAGESIZE"])
            self.page_size = int(result.stdout)
        except (AdbError, ValueError):
            self.page_size = 0

    def _set_build_fingerprint(self) -> None:
        if not self._adb:
            return
        try:
            result = self._adb.adb_command(
                ["shell", "getprop", "ro.build.fingerprint"])
            self.build_fingerprint = result.stdout.strip()
        except AdbError:
            self.build_fingerprint = ""

    def _record_device_details_in_sqlite_db(self):
        cursor = self.database_connection.cursor()
        cursor.execute(
            """\
            CREATE TABLE IF NOT EXISTS device_details (
                id INTEGER PRIMARY KEY,
                device_name TEXT NOT NULL,
                serial TEXT,
                page_size INTEGER,
                build_fingerprint TEXT
            )""")
        cursor.execute(
            """\
            INSERT INTO device_details (
                device_name,
                serial,
                page_size,
                build_fingerprint
            ) VALUES (?, ?, ?, ?)
            """,
            (
                self.device_name,
                self.serial,
                self.page_size,
                self.build_fingerprint,
            ),
        )
        self.database_connection.commit()

    def __str__(self) -> str:
        return (
            "AndroidDevice(\n"
            f"\tdevice_name={self.device_name},\n"
            f"\tserial='{self.serial}',\n"
            f"\tpage_size={self.page_size}, \n"
            f"\tbuild_fingerprint={self.build_fingerprint}, \n"
            f")"
        )

    def reboot(self) -> None:
        """Reboots the device."""
        if not self._adb:
            raise AdbError("Cannot reboot device, ADB not initialized.")
        print("Rebooting device and waiting 5 minutes...")
        self._adb.adb_command(["reboot"], check=False)
        self._adb.adb_command(["wait-for-device"], check=False)

    def record_logcat_logs(self) -> None:
        """Records logcat logs."""
        self._adb.root()
        self.post_boot_logcat_log = self._adb.adb_command(
            ["shell", "logcat", "-b", "kernel", "-d"], check=False,
        ).stdout

    def fetch_and_store_memory_data(self, run_config: RunConfig) -> None:
        """Runs specified workloads, system probes to generate results."""
        from runnables.runnable_data import RunnableData  # pylint: disable=import-outside-toplevel
        runnable_data = RunnableData()
        runnable_data.android_device = self
        runnable_data.run_config = run_config
        runnable_prompt_device_serial.process(runnable_data)
        runnable_apply_device_configs.process(runnable_data)

        # Circular reference errors caused by pylint's static analyzer.
        # pylint: disable=no-member
        print(
            TerminalPrintColors.BOLD,
            TerminalPrintColors.MAGENTA,
            f"\n\n├────────── Starting Analysis for device "
            f"{runnable_data.android_device_not_null.device_name} "
            f"({runnable_data.android_device_not_null.serial}) "
            f"(page size: {runnable_data.android_device_not_null.page_size:,}) "
            "─────────┤",
            TerminalPrintColors.RESET,
        )
        # pylint: enable=no-member

        for reboot_count in range(
            1,
            runnable_data.run_config_not_null.num_reboots + 1
        ):
            self.reboot()
            self.record_logcat_logs()
            time.sleep(300)

            runnable_data.reboot_count = reboot_count

            # After reboot, adb needs to be set to root again.
            if not self._adb:
                self._adb = Adb(self.serial)

            while True:
                try:
                    print("Running adb as root after reboot...")
                    self._adb.root()
                    break
                except AdbError as e:
                    print(f"Failed to set root after reboot: {e}")
                    # pylint: disable=bad-builtin
                    input(
                        "Please check device connection and "
                        "press Enter to retry..."
                    )

            print(
                f"  ├─ Reboot run #{reboot_count} of "
                f"{runnable_data.run_config_not_null.num_reboots}"
            )
            for workload in runnable_data.run_config_not_null.workloads:
                for probe_sample_count in range(
                    1,
                    runnable_data.run_config_not_null.runs_per_workload +1
                ):
                    runnable_data.workload = workload
                    runnable_data.probe_sample_count = probe_sample_count

                    system_probes_print_string = ",".join(
                        [
                            p.name for p in
                            runnable_data.run_config_not_null.system_probes
                        ]
                    )
                    print(
                        f"    └── [{system_probes_print_string}] + ",
                        f"{runnable_data.workload_not_null.name}",
                    )

                    # Run workload.
                    runnable_run_workload.process(runnable_data)

                    # Run system probes.
                    runnable_run_system_probe.process(runnable_data)

                    # Post process (typically just close the app).
                    runnable_run_workload.post_process(runnable_data)

                # end for _
            # end for workload