1#!/usr/bin/env python
2#
3# Copyright (C) 2018 The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#      http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17
18"""
19Utils for running unittests.
20"""
21
22import logging
23import os
24import os.path
25import struct
26import sys
27import unittest
28
29import common
30
31# Some test runner doesn't like outputs from stderr.
32logging.basicConfig(stream=sys.stdout)
33
34# Use ANDROID_BUILD_TOP as an indicator to tell if the needed tools (e.g.
35# avbtool, mke2fs) are available while running the tests, unless
36# FORCE_RUN_RELEASETOOLS is set to '1'. Not having the required vars means we
37# can't run the tests that require external tools.
38EXTERNAL_TOOLS_UNAVAILABLE = (
39    not os.environ.get('ANDROID_BUILD_TOP') and
40    os.environ.get('FORCE_RUN_RELEASETOOLS') != '1')
41
42
43def SkipIfExternalToolsUnavailable():
44  """Decorator function that allows skipping tests per tools availability."""
45  if EXTERNAL_TOOLS_UNAVAILABLE:
46    return unittest.skip('External tools unavailable')
47  return lambda func: func
48
49
50def get_testdata_dir():
51  """Returns the testdata dir, in relative to the script dir."""
52  # The script dir is the one we want, which could be different from pwd.
53  current_dir = os.path.dirname(os.path.realpath(__file__))
54  return os.path.join(current_dir, 'testdata')
55
56
57def get_search_path():
58  """Returns the search path that has 'framework/signapk.jar' under."""
59
60  def signapk_exists(path):
61    signapk_path = os.path.realpath(
62        os.path.join(path, 'framework', 'signapk.jar'))
63    return os.path.exists(signapk_path)
64
65  # Try with ANDROID_BUILD_TOP first.
66  full_path = os.path.realpath(os.path.join(
67      os.environ.get('ANDROID_BUILD_TOP', ''), 'out', 'host', 'linux-x86'))
68  if signapk_exists(full_path):
69    return full_path
70
71  # Otherwise try going with relative pathes.
72  current_dir = os.path.dirname(os.path.realpath(__file__))
73  for path in (
74      # In relative to 'build/make/tools/releasetools' in the Android source.
75      ['..'] * 4 + ['out', 'host', 'linux-x86'],
76      # Or running the script unpacked from otatools.zip.
77      ['..']):
78    full_path = os.path.realpath(os.path.join(current_dir, *path))
79    if signapk_exists(full_path):
80      return full_path
81  return None
82
83
84def construct_sparse_image(chunks):
85  """Returns a sparse image file constructed from the given chunks.
86
87  From system/core/libsparse/sparse_format.h.
88  typedef struct sparse_header {
89    __le32 magic;  // 0xed26ff3a
90    __le16 major_version;  // (0x1) - reject images with higher major versions
91    __le16 minor_version;  // (0x0) - allow images with higer minor versions
92    __le16 file_hdr_sz;  // 28 bytes for first revision of the file format
93    __le16 chunk_hdr_sz;  // 12 bytes for first revision of the file format
94    __le32 blk_sz;  // block size in bytes, must be a multiple of 4 (4096)
95    __le32 total_blks;  // total blocks in the non-sparse output image
96    __le32 total_chunks;  // total chunks in the sparse input image
97    __le32 image_checksum;  // CRC32 checksum of the original data, counting
98                            // "don't care" as 0. Standard 802.3 polynomial,
99                            // use a Public Domain table implementation
100  } sparse_header_t;
101
102  typedef struct chunk_header {
103    __le16 chunk_type;  // 0xCAC1 -> raw; 0xCAC2 -> fill;
104                        // 0xCAC3 -> don't care
105    __le16 reserved1;
106    __le32 chunk_sz;  // in blocks in output image
107    __le32 total_sz;  // in bytes of chunk input file including chunk header
108                      // and data
109  } chunk_header_t;
110
111  Args:
112    chunks: A list of chunks to be written. Each entry should be a tuple of
113        (chunk_type, block_number).
114
115  Returns:
116    Filename of the created sparse image.
117  """
118  SPARSE_HEADER_MAGIC = 0xED26FF3A
119  SPARSE_HEADER_FORMAT = "<I4H4I"
120  CHUNK_HEADER_FORMAT = "<2H2I"
121
122  sparse_image = common.MakeTempFile(prefix='sparse-', suffix='.img')
123  with open(sparse_image, 'wb') as fp:
124    fp.write(struct.pack(
125        SPARSE_HEADER_FORMAT, SPARSE_HEADER_MAGIC, 1, 0, 28, 12, 4096,
126        sum(chunk[1] for chunk in chunks),
127        len(chunks), 0))
128
129    for chunk in chunks:
130      data_size = 0
131      if chunk[0] == 0xCAC1:
132        data_size = 4096 * chunk[1]
133      elif chunk[0] == 0xCAC2:
134        data_size = 4
135      elif chunk[0] == 0xCAC3:
136        pass
137      else:
138        assert False, "Unsupported chunk type: {}".format(chunk[0])
139
140      fp.write(struct.pack(
141          CHUNK_HEADER_FORMAT, chunk[0], 0, chunk[1], data_size + 12))
142      if data_size != 0:
143        fp.write(os.urandom(data_size))
144
145  return sparse_image
146
147
148class MockScriptWriter(object):
149  """A class that mocks edify_generator.EdifyGenerator.
150
151  It simply pushes the incoming arguments onto script stack, which is to assert
152  the calls to EdifyGenerator functions.
153  """
154
155  def __init__(self, enable_comments=False):
156    self.lines = []
157    self.enable_comments = enable_comments
158
159  def Mount(self, *args):
160    self.lines.append(('Mount',) + args)
161
162  def AssertDevice(self, *args):
163    self.lines.append(('AssertDevice',) + args)
164
165  def AssertOemProperty(self, *args):
166    self.lines.append(('AssertOemProperty',) + args)
167
168  def AssertFingerprintOrThumbprint(self, *args):
169    self.lines.append(('AssertFingerprintOrThumbprint',) + args)
170
171  def AssertSomeFingerprint(self, *args):
172    self.lines.append(('AssertSomeFingerprint',) + args)
173
174  def AssertSomeThumbprint(self, *args):
175    self.lines.append(('AssertSomeThumbprint',) + args)
176
177  def Comment(self, comment):
178    if not self.enable_comments:
179      return
180    self.lines.append('# {}'.format(comment))
181
182  def AppendExtra(self, extra):
183    self.lines.append(extra)
184
185  def __str__(self):
186    return '\n'.join(self.lines)
187
188
189class ReleaseToolsTestCase(unittest.TestCase):
190  """A common base class for all the releasetools unittests."""
191
192  def tearDown(self):
193    common.Cleanup()
194
195
196if __name__ == '__main__':
197  testsuite = unittest.TestLoader().discover(
198      os.path.dirname(os.path.realpath(__file__)))
199  # atest needs a verbosity level of >= 2 to correctly parse the result.
200  unittest.TextTestRunner(verbosity=2).run(testsuite)
201