1#!/usr/bin/env python3
2#
3#   Copyright 2019 - 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
17import collections
18import os
19import tempfile
20
21from acts.test_utils.instrumentation.proto.gen import instrumentation_data_pb2
22
23DEFAULT_INST_LOG_DIR = 'instrument-logs'
24
25START_TIMESTAMP = 'start'
26END_TIMESTAMP = 'end'
27
28
29class ProtoParserError(Exception):
30    """Class for exceptions raised by the proto parser."""
31
32
33def pull_proto(ad, dest_dir, source_path=None):
34    """Pull latest instrumentation result proto from device.
35
36    Args:
37        ad: AndroidDevice object
38        dest_dir: Directory on the host where the proto will be sent
39        source_path: Path on the device where the proto is generated. If None,
40            pull the latest proto from DEFAULT_INST_PROTO_DIR.
41
42    Returns: Path to the retrieved proto file
43    """
44    if source_path:
45        filename = os.path.basename(source_path)
46    else:
47        default_full_proto_dir = os.path.join(
48            ad.external_storage_path, DEFAULT_INST_LOG_DIR)
49        filename = ad.adb.shell('ls %s -t | head -n1' % default_full_proto_dir)
50        if not filename:
51            raise ProtoParserError(
52                'No instrumentation result protos found at default location.')
53        source_path = os.path.join(default_full_proto_dir, filename)
54    ad.pull_files(source_path, dest_dir)
55    dest_path = os.path.join(dest_dir, filename)
56    if not os.path.exists(dest_path):
57        raise ProtoParserError(
58            'Failed to pull instrumentation result proto: %s -> %s'
59            % (source_path, dest_path))
60    return dest_path
61
62
63def get_session_from_local_file(proto_file):
64    """Get a instrumentation_data_pb2.Session object from a proto file on the
65    host.
66
67    Args:
68        proto_file: Path to the proto file (on host)
69
70    Returns: A instrumentation_data_pb2.Session
71    """
72    with open(proto_file, 'rb') as f:
73        return instrumentation_data_pb2.Session.FromString(f.read())
74
75
76def get_session_from_device(ad, proto_file=None):
77    """Get a instrumentation_data_pb2.Session object from a proto file on
78    device.
79
80    Args:
81        ad: AndroidDevice object
82        proto_file: Path to the proto file (on device). If None, defaults to
83            latest proto from DEFAULT_INST_PROTO_DIR.
84
85    Returns: A instrumentation_data_pb2.Session
86    """
87    with tempfile.TemporaryDirectory() as tmp_dir:
88        pulled_proto = pull_proto(ad, tmp_dir, proto_file)
89        return get_session_from_local_file(pulled_proto)
90
91
92def get_test_timestamps(session):
93    """Parse an instrumentation_data_pb2.Session to get the timestamps for each
94    test.
95
96    Args:
97        session: an instrumentation_data.Session object
98
99    Returns: a dict in the format
100        {
101            <test name> : (<begin_time>, <end_time>),
102            ...
103        }
104    """
105    timestamps = collections.defaultdict(dict)
106    for test_status in session.test_status:
107        entries = test_status.results.entries
108        # Timestamp entries have the key 'timestamp-message'
109        if any(entry.key == 'timestamps-message' for entry in entries):
110            test_name = None
111            timestamp = None
112            timestamp_type = None
113            for entry in entries:
114                if entry.key == 'test':
115                    test_name = entry.value_string
116                if entry.key == 'timestamp':
117                    timestamp = entry.value_long
118                if entry.key == 'start-timestamp':
119                    timestamp_type = START_TIMESTAMP
120                if entry.key == 'end-timestamp':
121                    timestamp_type = END_TIMESTAMP
122            if test_name and timestamp and timestamp_type:
123                timestamps[test_name][timestamp_type] = timestamp
124    return timestamps
125