1# Copyright 2018 Google Inc.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Module for the snippet management service."""
15from mobly.controllers.android_device_lib import errors
16from mobly.controllers.android_device_lib import snippet_client_v2
17from mobly.controllers.android_device_lib.services import base_service
18
19MISSING_SNIPPET_CLIENT_MSG = 'No snippet client is registered with name "%s".'
20
21
22class Error(errors.ServiceError):
23  """Root error type for snippet management service."""
24
25  SERVICE_TYPE = 'SnippetManagementService'
26
27
28class SnippetManagementService(base_service.BaseService):
29  """Management service of snippet clients.
30
31  This service manages all the snippet clients associated with an Android
32  device.
33  """
34
35  def __init__(self, device, configs=None):
36    del configs  # Unused param.
37    self._device = device
38    self._is_alive = False
39    self._snippet_clients = {}
40    super().__init__(device)
41
42  @property
43  def is_alive(self):
44    """True if any client is running, False otherwise."""
45    return any([client.is_alive for client in self._snippet_clients.values()])
46
47  def get_snippet_client(self, name):
48    """Gets the snippet client managed under a given name.
49
50    Args:
51      name: string, the name of the snippet client under management.
52
53    Returns:
54      SnippetClient.
55    """
56    if name in self._snippet_clients:
57      return self._snippet_clients[name]
58
59  def add_snippet_client(self, name, package, config=None):
60    """Adds a snippet client to the management.
61
62    Args:
63      name: string, the attribute name to which to attach the snippet
64        client. E.g. `name='maps'` attaches the snippet client to
65        `ad.maps`.
66      package: string, the package name of the snippet apk to connect to.
67      config: snippet_client_v2.Config, the configuration object for
68        controlling the snippet behaviors. See the docstring of the `Config`
69        class for supported configurations.
70
71    Raises:
72      Error, if a duplicated name or package is passed in.
73    """
74    # Should not load snippet with the same name more than once.
75    if name in self._snippet_clients:
76      raise Error(
77          self,
78          'Name "%s" is already registered with package "%s", it cannot '
79          'be used again.' % (name, self._snippet_clients[name].client.package),
80      )
81    # Should not load the same snippet package more than once.
82    for snippet_name, client in self._snippet_clients.items():
83      if package == client.package:
84        raise Error(
85            self,
86            'Snippet package "%s" has already been loaded under name "%s".'
87            % (package, snippet_name),
88        )
89
90    client = snippet_client_v2.SnippetClientV2(
91        package=package,
92        ad=self._device,
93        config=config,
94    )
95    client.initialize()
96    self._snippet_clients[name] = client
97
98  def remove_snippet_client(self, name):
99    """Removes a snippet client from management.
100
101    Args:
102      name: string, the name of the snippet client to remove.
103
104    Raises:
105      Error: if no snippet client is managed under the specified name.
106    """
107    if name not in self._snippet_clients:
108      raise Error(self._device, MISSING_SNIPPET_CLIENT_MSG % name)
109    client = self._snippet_clients.pop(name)
110    client.stop()
111
112  def start(self):
113    """Starts all the snippet clients under management."""
114    for client in self._snippet_clients.values():
115      if not client.is_alive:
116        self._device.log.debug('Starting SnippetClient<%s>.', client.package)
117        client.initialize()
118      else:
119        self._device.log.debug(
120            'Not startng SnippetClient<%s> because it is already alive.',
121            client.package,
122        )
123
124  def stop(self):
125    """Stops all the snippet clients under management."""
126    for client in self._snippet_clients.values():
127      if client.is_alive:
128        self._device.log.debug('Stopping SnippetClient<%s>.', client.package)
129        client.stop()
130      else:
131        self._device.log.debug(
132            'Not stopping SnippetClient<%s> because it is not alive.',
133            client.package,
134        )
135
136  def pause(self):
137    """Pauses all the snippet clients under management.
138
139    This clears the host port of a client because a new port will be
140    allocated in `resume`.
141    """
142    for client in self._snippet_clients.values():
143      self._device.log.debug('Pausing SnippetClient<%s>.', client.package)
144      client.close_connection()
145
146  def resume(self):
147    """Resumes all paused snippet clients."""
148    for client in self._snippet_clients.values():
149      if not client.is_alive:
150        self._device.log.debug('Resuming SnippetClient<%s>.', client.package)
151        client.restore_server_connection()
152      else:
153        self._device.log.debug(
154            'Not resuming SnippetClient<%s>.', client.package
155        )
156
157  def __getattr__(self, name):
158    client = self.get_snippet_client(name)
159    if client:
160      return client
161    return self.__getattribute__(name)
162