packages/antlion/controllers/fuchsia_lib/ssh.py

#!/usr/bin/env python3
#
# Copyright 2022 The Fuchsia 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 time

from antlion.capabilities.ssh import SSHError, SSHProvider

DEFAULT_SSH_USER: str = "fuchsia"
DEFAULT_SSH_PRIVATE_KEY: str = "~/.ssh/fuchsia_ed25519"
# The default package repository for all components.
FUCHSIA_PACKAGE_REPO_NAME = "fuchsia.com"


class FuchsiaSSHProvider(SSHProvider):
    """Device-specific provider for SSH clients."""

    def start_v1_component(
        self,
        component: str,
        timeout_sec: int = 5,
        repo: str = FUCHSIA_PACKAGE_REPO_NAME,
    ) -> None:
        """Start a CFv1 component in the background.

        Args:
            component: Name of the component without ".cmx".
            timeout_sec: Seconds to wait for the process to show up in 'ps'.
            repo: Default package repository for all components.

        Raises:
            TimeoutError: when the component doesn't launch within timeout_sec
        """
        # The "run -d" command will hang when executed without a pseudo-tty
        # allocated.
        self.run(
            f"run -d fuchsia-pkg://{repo}/{component}#meta/{component}.cmx",
            force_tty=True,
        )

        timeout = time.perf_counter() + timeout_sec
        while True:
            ps_cmd = self.run("ps")
            if f"{component}.cmx" in ps_cmd.stdout:
                return
            if time.perf_counter() > timeout:
                raise TimeoutError(
                    f'Failed to start "{component}.cmx" after {timeout_sec}s'
                )

    def stop_component(self, component: str, is_cfv2_component: bool = False) -> None:
        """Stop all instances of a CFv1 or CFv2 component.

        Args:
            component: Name of the component without suffix("cm" or "cmx").
            is_cfv2_component: Determines the component suffix to use.
        """
        suffix = "cm" if is_cfv2_component else "cmx"

        try:
            self.run(f"killall {component}.{suffix}")
            self.log.info(f"Stopped component: {component}.{suffix}")
        except SSHError as e:
            if "no tasks found" in e.result.stderr:
                self.log.debug(f"Could not find component: {component}.{suffix}")
                return
            raise e