devlib initial commit.

devlib/instrument/netstats/__init__.py

@@ -0,0 +1,135 @@
+import os
+import re
+import csv
+import tempfile
+from datetime import datetime
+from collections import defaultdict
+from itertools import izip_longest
+
+from devlib.instrument import Instrument, MeasurementsCsv, CONTINUOUS
+from devlib.exception import TargetError, HostError
+from devlib.utils.android import ApkInfo
+
+
+THIS_DIR = os.path.dirname(__file__)
+
+NETSTAT_REGEX = re.compile(r'I/(?P<tag>netstats-\d+)\(\s*\d*\): (?P<ts>\d+) '
+                           r'"(?P<package>[^"]+)" TX: (?P<tx>\S+) RX: (?P<rx>\S+)')
+
+
+def extract_netstats(filepath, tag=None):
+    netstats = []
+    with open(filepath) as fh:
+        for line in fh:
+            match = NETSTAT_REGEX.search(line)
+            if not match:
+                continue
+            if tag and match.group('tag') != tag:
+                continue
+            netstats.append((match.group('tag'),
+                             match.group('ts'),
+                             match.group('package'),
+                             match.group('tx'),
+                             match.group('rx')))
+    return netstats
+
+
+def netstats_to_measurements(netstats):
+    measurements = defaultdict(list)
+    for row in netstats:
+        tag, ts, package, tx, rx = row  # pylint: disable=unused-variable
+        measurements[package + '_tx'].append(tx)
+        measurements[package + '_rx'].append(rx)
+    return measurements
+
+
+def write_measurements_csv(measurements, filepath):
+    headers = sorted(measurements.keys())
+    columns = [measurements[h] for h in headers]
+    with open(filepath, 'wb') as wfh:
+        writer = csv.writer(wfh)
+        writer.writerow(headers)
+        writer.writerows(izip_longest(*columns))
+
+
+class NetstatsInstrument(Instrument):
+
+    mode = CONTINUOUS
+
+    def __init__(self, target, apk=None, service='.TrafficMetricsService'):
+        """
+        Additional paramerter:
+
+        :apk: Path to the APK file that contains ``com.arm.devlab.netstats``
+              package. If not specified, it will be assumed that an APK with
+              name "netstats.apk" is located in the same directory as the
+              Python module for the instrument.
+        :service: Name of the service to be launched. This service must be
+                  present in the APK.
+
+        """
+        if target.os != 'android':
+            raise TargetError('netstats insturment only supports Android targets')
+        if apk is None:
+            apk = os.path.join(THIS_DIR, 'netstats.apk')
+        if not os.path.isfile(apk):
+            raise HostError('APK for netstats instrument does not exist ({})'.format(apk))
+        super(NetstatsInstrument, self).__init__(target)
+        self.apk = apk
+        self.package = ApkInfo(self.apk).package
+        self.service = service
+        self.tag = None
+        self.command = None
+        self.stop_command = 'am kill {}'.format(self.package)
+
+        for package in self.target.list_packages():
+            self.add_channel(package, 'tx')
+            self.add_channel(package, 'rx')
+
+    def setup(self, force=False, *args, **kwargs):
+        if self.target.package_is_installed(self.package):
+            if force:
+                self.logger.debug('Re-installing {} (forced)'.format(self.package))
+                self.target.uninstall_package(self.package)
+                self.target.install(self.apk)
+            else:
+                self.logger.debug('{} already present on target'.format(self.package))
+        else:
+            self.logger.debug('Deploying {} to target'.format(self.package))
+            self.target.install(self.apk)
+
+    def reset(self, sites=None, kinds=None, period=None):  # pylint: disable=arguments-differ
+        super(NetstatsInstrument, self).reset(sites, kinds)
+        period_arg, packages_arg = '', ''
+        self.tag = 'netstats-{}'.format(datetime.now().strftime('%Y%m%d%H%M%s'))
+        tag_arg = ' --es tag {}'.format(self.tag)
+        if sites:
+            packages_arg = ' --es packages {}'.format(','.join(sites))
+        if period:
+            period_arg = ' --ei period {}'.format(period)
+        self.command = 'am startservice{}{}{} {}/{}'.format(tag_arg,
+                                                            period_arg,
+                                                            packages_arg,
+                                                            self.package,
+                                                            self.service)
+        self.target.execute(self.stop_command)  # ensure the service is not running.
+
+    def start(self):
+        if self.command is None:
+            raise RuntimeError('reset() must be called before start()')
+        self.target.execute(self.command)
+
+    def stop(self):
+        self.target.execute(self.stop_command)
+
+    def get_data(self, outfile):
+        raw_log_file = tempfile.mktemp()
+        self.target.dump_logcat(raw_log_file)
+        data = extract_netstats(raw_log_file)
+        measurements = netstats_to_measurements(data)
+        write_measurements_csv(measurements, outfile)
+        os.remove(raw_log_file)
+        return MeasurementsCsv(outfile, self.active_channels)
+
+    def teardown(self):
+        self.target.uninstall_package(self.package)