xref: /aosp_15_r20/external/autotest/server/hosts/tls_client/fake_omaha.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7)
1# Lint as: python2, python3
2# Copyright 2021 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5"""Communication with the TLS FakeOmaha Service."""
6
7import logging
8
9import common
10
11from autotest_lib.server.hosts.tls_client import autotest_common_pb2
12from autotest_lib.client.common_lib import error
13
14PAYLOAD_TYPE = {
15        'TYPE_UNSPECIFIED':
16        autotest_common_pb2.FakeOmaha.Payload.TYPE_UNSPECIFIED,
17        'FULL': autotest_common_pb2.FakeOmaha.Payload.FULL,
18        'DELTA': autotest_common_pb2.FakeOmaha.Payload.DELTA
19}
20
21
22class TLSFakeOmaha():
23    """Object for sending commands to a host, and getting the response."""
24
25    def __init__(self, tlsconnection):
26        """Configure the grpc channel."""
27        if tlsconnection.alive:
28            self.stub = tlsconnection.stub
29        else:
30            raise error.TLSConnectionError(
31                    "TLS connection is not alive when try to creating"
32                    " FakeOmaha client.")
33
34        self.tlsconnection = tlsconnection
35        self.fo_name = None
36
37    def _make_payloads(self, payloads):
38        """Serialize and return the list of payloads."""
39        serialized_payloads = []
40        for payload in payloads:
41            serialized_payloads.append(
42                    autotest_common_pb2.FakeOmaha.Payload(
43                            id=payload['payload_id'],
44                            type=PAYLOAD_TYPE[payload['payload_type']]))
45
46        return serialized_payloads
47
48    def start_omaha(self,
49                    hostname,
50                    target_build,
51                    payloads,
52                    exposed_via_proxy=False,
53                    critical_update=False,
54                    return_noupdate_starting=0):
55        """Serialize and send the cmd to the TLS service.
56
57        @param hostname: hostname of dut. Normally 'hostname' or 'self.hostname'
58        @param target_build: full target build for the update. Example:
59
60        @param payloads: list of the payloads in the format:
61            [{'payload_id': <id>, 'payload_type': <type>}]
62            example:
63                [{'payload_id': 'ROOTFS', 'payload_type': 'FULL'},]
64        @param exposed_via_proxy: bool indicates that the fake Omaha service is
65            exposed to a DUT via a proxy server, instead of exposing to the DUT
66                directly.
67        @param critical_update:bool instructs the fake Omaha created that the
68            update is critical.
69        @param return_noupdate_starting: int indicates from which update check
70            to start returning noupdate.
71
72        @returns: the omaha_url
73        """
74        payload = self._make_payloads(payloads)
75
76        target_build = autotest_common_pb2.ChromeOsImage(
77                gs_path_prefix=target_build)
78        fake_omaha = autotest_common_pb2.FakeOmaha(
79                dut=hostname,
80                target_build=target_build,
81                payloads=payload,
82                exposed_via_proxy=exposed_via_proxy,
83                critical_update=critical_update,
84                return_noupdate_starting=return_noupdate_starting)
85
86        req = autotest_common_pb2.CreateFakeOmahaRequest(fake_omaha=fake_omaha)
87
88        try:
89            result = self.stub.CreateFakeOmaha(req)
90            self.fo_name = result.name
91            return result.omaha_url
92        except Exception as e:
93            logging.error("TLS FakeOmaha Debug String: %s",
94                          e.debug_error_string())
95            raise error.TestError(
96                    "Could not start FakeOmaha Server because %s", e.details())
97
98    def stop_omaha(self):
99        """Delete the running Omaha Service."""
100        if not self.fo_name:
101            raise error.TestWarn(
102                    "No FakeOmaha name specified, has it been started?")
103        req = autotest_common_pb2.DeleteFakeOmahaRequest(name=self.fo_name)
104        try:
105            self.stub.DeleteFakeOmaha(req)
106        except Exception as e:
107            raise error.TestWarn("Could not stop FakeOmaha Server because %s",
108                                 e.details())
109