1import argparse
2from importlib import resources
3import os
4import subprocess
5import sys
6import tempfile
7import time
8import config_pb2
9import google.protobuf.text_format as text_format
10
11kwargs = {
12    "text": True,
13    "shell": True,
14    "capture_output": True,
15    "check": True,
16}
17
18
19def adb_root():
20  subprocess.run("adb root", **kwargs)
21
22
23def create_and_push_config_proto(name="test_slog"):
24  with resources.open_text("test", f"{name}.textproto") as textproto:
25    message = text_format.Parse(
26        textproto.read(), config_pb2.UprobestatsConfig()
27    )
28    textproto.close()
29
30    with tempfile.NamedTemporaryFile() as temp:
31      temp.write(message.SerializeToString())
32      temp.flush()
33      print(f"creating {name}")
34      config_cmd = (
35          f"adb push {temp.name} /data/misc/uprobestats-configs/{name}.proto"
36      )
37      subprocess.run(config_cmd, **kwargs)
38
39
40def clear_logcat():
41  print("clearing logcat")
42  subprocess.run("adb logcat -c", **kwargs)
43
44
45def start_uprobestats(name="test_slog"):
46  print("starting uprobestats")
47  subprocess.run(
48      f"adb shell setprop uprobestats.start_with_config {name}.proto", **kwargs
49  )
50
51
52def get_ring_buffer_values():
53  time.sleep(10)
54  lines = subprocess.run(
55      'adb logcat -d | grep "ringbuf result callback. value:"', **kwargs
56  ).stdout.splitlines()
57  for line in lines:
58    print(line)
59
60
61def get_ring_buffer_size():
62  # print('generating log messages and fetching ring buffer size')
63  # subprocess.run('adb shell killall com.google.android.apps.photos', shell=True)
64  # time.sleep(2)
65  # subprocess.run(
66  #    'adb shell monkey -p com.google.android.apps.photos -c'
67  #    ' android.intent.category.LAUNCHER 1',
68  #    **kwargs,
69  # )
70  ring_buffer_size = (
71      subprocess.run(
72          'adb logcat -d | grep "uprobestats: ring buffer size"', **kwargs
73      )
74      .stdout.splitlines()[0]
75      .split(":")[4]
76      .split(" ")[1]
77      .strip()
78  )
79
80  print(f"ring buffer size: {ring_buffer_size}")
81  return int(ring_buffer_size, 0)
82
83
84if __name__ == "__main__":
85  parser = argparse.ArgumentParser(
86      "Drops a uprobestats config over adb and checks start it. Optionally acts"
87      " as a test by checking logcat output"
88  )
89  parser.add_argument(
90      "-n",
91      "--name",
92      type=str,
93      default="test_slog",
94      help="Name of the config file e.g. test_slog",
95  )
96  parser.add_argument(
97      "-t",
98      "--test",
99      default=False,
100      action="store_true",
101      help=(
102          "Run a test, meaning start uprobestats and then look for particular"
103          " logcat output to exit successfully"
104      ),
105  )
106  args = parser.parse_args()
107
108  adb_root()
109  create_and_push_config_proto(args.name)
110
111  if not args.test:
112    start_uprobestats(args.name)
113    sys.exit(0)
114
115  clear_logcat()
116  start_uprobestats(args.name)
117  time.sleep(60)
118  ring_buf = get_ring_buffer_size()
119  get_ring_buffer_values()
120  if ring_buf > 0:
121    sys.exit(0)
122  else:
123    raise SystemExit("Did not find a ring buffer greater than zero")
124