python/test/api_integrationtest.py

#!/usr/bin/env python3
# Copyright (C) 2020 The Android Open Source Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import io
import os
import sqlite3
import subprocess
import tempfile
import unittest
from typing import Optional
from unittest import mock

import pandas as pd

from perfetto.batch_trace_processor.api import BatchTraceProcessor
from perfetto.batch_trace_processor.api import BatchTraceProcessorConfig
from perfetto.batch_trace_processor.api import FailureHandling
from perfetto.batch_trace_processor.api import Metadata
from perfetto.batch_trace_processor.api import TraceListReference
from perfetto.trace_processor.protos import ProtoFactory
from perfetto.trace_processor.api import PLATFORM_DELEGATE
from perfetto.trace_processor.api import TraceProcessor
from perfetto.trace_processor.api import TraceProcessorException
from perfetto.trace_processor.api import TraceProcessorConfig
from perfetto.trace_processor.api import TraceReference
from perfetto.trace_uri_resolver.resolver import TraceUriResolver
from perfetto.trace_uri_resolver.path import PathUriResolver


class SimpleResolver(TraceUriResolver):
  PREFIX = 'simple'

  def __init__(self, path, skip_resolve_file=False):
    self.path = path
    self.file = open(example_android_trace_path(), 'rb')
    self.skip_resolve_file = skip_resolve_file

  def file_gen(self):
    with open(example_android_trace_path(), 'rb') as f:
      yield f.read()

  def resolve(self):
    res = [
        TraceUriResolver.Result(
            self.file_gen(), metadata={'source': 'generator'}),
        TraceUriResolver.Result(
            example_android_trace_path(), metadata={'source': 'path'}),
    ]
    if not self.skip_resolve_file:
      res.extend([
          TraceUriResolver.Result(
              PathUriResolver(example_android_trace_path()),
              metadata={'source': 'path_resolver'}),
          TraceUriResolver.Result(self.file, metadata={'source': 'file'}),
      ])
    return res


class RecursiveResolver(SimpleResolver):
  PREFIX = 'recursive'

  def __init__(self, path, skip_resolve_file):
    super().__init__(path=path, skip_resolve_file=skip_resolve_file)

  def resolve(self):
    srf = self.skip_resolve_file
    return [
        TraceUriResolver.Result(
            self.file_gen(), metadata={'source': 'recursive_gen'}),
        TraceUriResolver.Result(
            f'simple:path={self.path};skip_resolve_file={srf}',
            metadata={
                'source': 'recursive_path',
                'root_source': 'recursive_path'
            }),
        TraceUriResolver.Result(
            SimpleResolver(
                path=self.path, skip_resolve_file=self.skip_resolve_file),
            metadata={
                'source': 'recursive_obj',
                'root_source': 'recursive_obj'
            }),
    ]


class SimpleObserver(BatchTraceProcessor.Observer):

  def __init__(self):
    self.execution_times = []

  def trace_processed(self, metadata: Metadata, execution_time_seconds: float):
    self.execution_times.append(execution_time_seconds)


def create_batch_tp(
    traces: TraceListReference,
    load_failure_handling: FailureHandling = FailureHandling.RAISE_EXCEPTION,
    execute_failure_handling: FailureHandling = FailureHandling.RAISE_EXCEPTION,
    observer: Optional[BatchTraceProcessor.Observer] = None):
  registry = PLATFORM_DELEGATE().default_resolver_registry()
  registry.register(SimpleResolver)
  registry.register(RecursiveResolver)
  config = BatchTraceProcessorConfig(
      load_failure_handling=load_failure_handling,
      execute_failure_handling=execute_failure_handling,
      tp_config=TraceProcessorConfig(
          bin_path=os.environ["SHELL_PATH"], resolver_registry=registry))
  return BatchTraceProcessor(traces=traces, config=config, observer=observer)


def create_tp(trace: TraceReference):
  return TraceProcessor(
      trace=trace,
      config=TraceProcessorConfig(bin_path=os.environ["SHELL_PATH"]))


def example_android_trace_path():
  return os.path.join(os.environ["ROOT_DIR"], 'test', 'data',
                      'example_android_trace_30s.pb')


class TestApi(unittest.TestCase):

  def test_invalid_trace(self):
    f = io.BytesIO(b'<foo></foo>')
    with self.assertRaises(TraceProcessorException):
      _ = create_tp(trace=f)

  def test_runtime_error(self):
    # We emulate a situation when TP returns an error by passing the --version
    # flag. This makes TP output version information and exit, instead of
    # starting an http server.
    config = TraceProcessorConfig(
        bin_path=os.environ["SHELL_PATH"], extra_flags=["--version"])
    with self.assertRaisesRegex(
        TraceProcessorException,
        expected_regex='.*Trace Processor RPC API version:.*'):
      TraceProcessor(trace=io.BytesIO(b''), config=config)

  def test_close_stops_shell(self):
    tp = create_tp(trace=example_android_trace_path())
    shell = tp.subprocess
    self.assertIsNone(shell.poll())
    tp.close()
    # wait() raises TimeoutExpired if the shell is still running.
    shell.wait(timeout=10)
    # A second close() must be a no-op.
    tp.close()

  def test_trace_path(self):
    # Get path to trace_processor_shell and construct TraceProcessor
    tp = create_tp(trace=example_android_trace_path())
    qr_iterator = tp.query('select * from slice limit 10')
    dur_result = [
        178646, 119740, 58073, 155000, 173177, 20209377, 3589167, 90104, 275312,
        65313
    ]

    for num, row in enumerate(qr_iterator):
      self.assertEqual(row.dur, dur_result[num])

    # Test the batching logic by issuing a large query and ensuring we receive
    # all rows, not just a truncated subset.
    qr_iterator = tp.query('select count(*) as cnt from slice')
    expected_count = next(qr_iterator).cnt
    self.assertGreater(expected_count, 0)

    qr_iterator = tp.query('select * from slice')
    count = sum(1 for _ in qr_iterator)
    self.assertEqual(count, expected_count)

    tp.close()

  def test_trace_byteio(self):
    f = io.BytesIO(
        b'\n(\n&\x08\x00\x12\x12\x08\x01\x10\xc8\x01\x1a\x0b\x12\t'
        b'B|200|foo\x12\x0e\x08\x02\x10\xc8\x01\x1a\x07\x12\x05E|200')
    with create_tp(trace=f) as tp:
      qr_iterator = tp.query('select * from slice limit 10')
      res = list(qr_iterator)

      self.assertEqual(len(res), 1)

      row = res[0]
      self.assertEqual(row.ts, 1)
      self.assertEqual(row.dur, 1)
      self.assertEqual(row.name, 'foo')

  def test_trace_file(self):
    with open(example_android_trace_path(), 'rb') as file:
      with create_tp(trace=file) as tp:
        qr_iterator = tp.query('select * from slice limit 10')
        dur_result = [
            178646, 119740, 58073, 155000, 173177, 20209377, 3589167, 90104,
            275312, 65313
        ]

        for num, row in enumerate(qr_iterator):
          self.assertEqual(row.dur, dur_result[num])

  def test_trace_generator(self):

    def reader_generator():
      with open(example_android_trace_path(), 'rb') as file:
        yield file.read(1024)

    with create_tp(trace=reader_generator()) as tp:
      qr_iterator = tp.query('select * from slice limit 10')
      dur_result = [
          178646, 119740, 58073, 155000, 173177, 20209377, 3589167, 90104,
          275312, 65313
      ]

      for num, row in enumerate(qr_iterator):
        self.assertEqual(row.dur, dur_result[num])

  def test_simple_resolver(self):
    sample_path = example_android_trace_path()
    dur = [178646, 178646, 178646, 178646]
    source = ['generator', 'path', 'path_resolver', 'file']

    # Only path and path_resolver will resolve to PathUriResolver so those will have the _path added
    # to their metadata
    path = [None, sample_path, sample_path, None]

    expected = pd.DataFrame(
        list(zip(dur, source, path)), columns=['dur', 'source', '_path'])

    with create_batch_tp(
        traces='simple:path={}'.format(example_android_trace_path())) as btp:
      df = btp.query_and_flatten('select dur from slice limit 1')
      pd.testing.assert_frame_equal(df, expected, check_dtype=False)

    with create_batch_tp(
        traces=SimpleResolver(path=example_android_trace_path())) as btp:
      df = btp.query_and_flatten('select dur from slice limit 1')
      pd.testing.assert_frame_equal(df, expected, check_dtype=False)

  def test_query_timing(self):
    observer = SimpleObserver()
    with create_batch_tp(
        traces='simple:path={}'.format(example_android_trace_path()),
        observer=observer) as btp:
      btp.query_and_flatten('select dur from slice limit 1')
      self.assertTrue(
          all([x > 0 for x in observer.execution_times]),
          'Running time should be positive')

  def test_recursive_resolver(self):
    dur = [
        178646, 178646, 178646, 178646, 178646, 178646, 178646, 178646, 178646
    ]
    source = ['recursive_gen', 'generator', 'path', 'generator', 'path']
    root_source = [
        float('nan'), 'recursive_path', 'recursive_path', 'recursive_obj',
        'recursive_obj'
    ]
    sample_path = example_android_trace_path()
    path = [None, None, sample_path, None, sample_path]
    expected = pd.DataFrame(
        list(zip(dur, source, root_source, path)),
        columns=['dur', 'source', 'root_source', '_path'])

    uri = 'recursive:path={};skip_resolve_file=true'.format(
        example_android_trace_path())
    with create_batch_tp(traces=uri) as btp:
      df = btp.query_and_flatten('select dur from slice limit 1')
      pd.testing.assert_frame_equal(df, expected, check_dtype=False)

    with create_batch_tp(
        traces=RecursiveResolver(
            path=example_android_trace_path(), skip_resolve_file=True)) as btp:
      df = btp.query_and_flatten('select dur from slice limit 1')
      pd.testing.assert_frame_equal(df, expected, check_dtype=False)

  def test_btp_load_failure(self):
    f = io.BytesIO(b'<foo></foo>')
    with self.assertRaises(TraceProcessorException):
      _ = create_batch_tp(traces=f)

  def test_btp_load_failure_increment_stat(self):
    f = io.BytesIO(b'<foo></foo>')
    btp = create_batch_tp(
        traces=f, load_failure_handling=FailureHandling.INCREMENT_STAT)
    self.assertEqual(btp.stats().load_failures, 1)

  def test_btp_query_failure(self):
    btp = create_batch_tp(traces=example_android_trace_path())
    with self.assertRaises(TraceProcessorException):
      _ = btp.query('select * from sl')

  def test_btp_query_failure_increment_stat(self):
    btp = create_batch_tp(
        traces=example_android_trace_path(),
        execute_failure_handling=FailureHandling.INCREMENT_STAT)
    _ = btp.query('select * from sl')
    self.assertEqual(btp.stats().execute_failures, 1)

  def test_btp_query_failure_message(self):
    btp = create_batch_tp(
        traces='simple:path={}'.format(example_android_trace_path()))
    with self.assertRaisesRegex(
        TraceProcessorException, expected_regex='.*source.*generator.*'):
      _ = btp.query('select * from sl')

  def test_extra_flags(self):
    with tempfile.TemporaryDirectory() as temp_dir:
      test_package_dir = os.path.join(temp_dir, 'ext')
      os.makedirs(test_package_dir)
      test_module = os.path.join(test_package_dir, 'module.sql')
      with open(test_module, 'w') as f:
        f.write('CREATE TABLE test_table AS SELECT 123 AS test_value\n')
      config = TraceProcessorConfig(
          bin_path=os.environ["SHELL_PATH"],
          extra_flags=['--add-sql-package', test_package_dir])
      with TraceProcessor(trace=io.BytesIO(b''), config=config) as tp:
        qr_iterator = tp.query(
            'SELECT IMPORT("ext.module"); SELECT test_value FROM test_table')
        self.assertEqual(next(qr_iterator).test_value, 123)

  def test_sql_file_access(self):
    with tempfile.TemporaryDirectory() as temp_dir:
      path = os.path.join(temp_dir, 'file_write')
      sql_path = path.replace("'", "''")

      config = TraceProcessorConfig(bin_path=os.environ["SHELL_PATH"])
      with TraceProcessor(trace=io.BytesIO(b''), config=config) as tp:
        with self.assertRaisesRegex(TraceProcessorException,
                                    'File I/O is disabled'):
          tp.query(f"SELECT __intrinsic_file_write('{sql_path}', X'000102FF')")

      config = TraceProcessorConfig(
          bin_path=os.environ["SHELL_PATH"], enable_sql_file_access=True)
      with TraceProcessor(trace=io.BytesIO(b''), config=config) as tp:
        row = next(
            tp.query(
                f"SELECT __intrinsic_file_write('{sql_path}', X'000102FF') AS size"
            ))
        self.assertEqual(row.size, 4)

      with open(path, 'rb') as f:
        self.assertEqual(f.read(), b'\x00\x01\x02\xff')

  def test_remote_sql_file_access_must_be_enabled_by_server(self):
    config = TraceProcessorConfig(enable_sql_file_access=True)
    with self.assertRaisesRegex(TraceProcessorException,
                                'server must be started with'):
      TraceProcessor(addr='localhost:1', config=config)

  def test_sqlite_export(self):
    with tempfile.TemporaryDirectory() as temp_dir:
      path = os.path.join(temp_dir, 'export.db')

      config = TraceProcessorConfig(bin_path=os.environ["SHELL_PATH"])
      with TraceProcessor(
          trace=example_android_trace_path(), config=config) as tp:
        tp.export(path, 'sqlite')

      con = sqlite3.connect(path)
      try:
        table_count = con.execute(
            "SELECT COUNT(*) FROM sqlite_master WHERE type='table'").fetchone(
            )[0]
        self.assertGreater(table_count, 0)
        self.assertGreater(
            con.execute('SELECT COUNT(*) FROM slice').fetchone()[0], 0)
      finally:
        con.close()

  def test_add_sql_packages(self):
    with tempfile.TemporaryDirectory() as temp_dir:
      # Create a directory structure for the package. The root of the
      # package is |temp_dir| and we are creating the file foo/bar.sql
      # inside it.
      test_package_dir = os.path.join(temp_dir, 'foo')
      os.makedirs(test_package_dir)
      test_module = os.path.join(test_package_dir, 'bar.sql')
      with open(test_module, 'w') as f:
        f.write(
            'CREATE PERFETTO TABLE test_sql_module_foo AS SELECT 1 AS value;\n')

      # Create another directory to add to the package to test that multiple
      # packages can be added.
      test_package_dir_2 = os.path.join(temp_dir, 'baz')
      os.makedirs(test_package_dir_2)
      test_module_2 = os.path.join(test_package_dir_2, 'qux.sql')
      with open(test_module_2, 'w') as f:
        f.write(
            'CREATE PERFETTO TABLE test_sql_module_foo_2 AS SELECT 2 AS value;\n'
        )

      config = TraceProcessorConfig(
          bin_path=os.environ["SHELL_PATH"],
          add_sql_packages=[test_package_dir, test_package_dir_2],
      )
      with TraceProcessor(trace=io.BytesIO(b''), config=config) as tp:
        qr_iterator = tp.query('INCLUDE PERFETTO MODULE foo.bar; '
                               'INCLUDE PERFETTO MODULE baz.qux; '
                               'SELECT value FROM test_sql_module_foo')
        self.assertEqual(next(qr_iterator).value, 1)

  def test_trace_summary_failure(self):
    tp = create_tp(trace=example_android_trace_path())
    with self.assertRaises(TraceProcessorException):
      _ = tp.trace_summary(['foo'], ['bar, baz'])
    tp.close()

  def test_trace_summary_success(self):
    metric_spec = """metric_spec: {
        id: "memory_per_process"
        value: "dur"
        query: {
          simple_slices {
            process_name_glob: "ab*"
          }
        }
      }
      """

    tp = create_tp(trace=example_android_trace_path())
    trace_summary = tp.trace_summary([metric_spec], ['memory_per_process'])
    self.assertEqual(trace_summary.metric_bundles[0].specs[0].id,
                     'memory_per_process')
    tp.close()

  def test_trace_summary_success_multiple_metrics(self):
    metric_spec_1 = """metric_spec: {
        id: "metric_one"
        value: "dur"
        query: {
          simple_slices {
            process_name_glob: "ab*"
          }
        }
      }
      """
    metric_spec_2 = """metric_spec: {
        id: "metric_two"
        value: "ts"
        query: {
          simple_slices {
            process_name_glob: "cd*"
          }
        }
      }
      """
    tp = create_tp(trace=example_android_trace_path())
    trace_summary = tp.trace_summary([metric_spec_1, metric_spec_2],
                                     ['metric_one', 'metric_two'])
    self.assertEqual(len(trace_summary.metric_bundles), 2)
    self.assertIn(trace_summary.metric_bundles[0].specs[0].id,
                  ['metric_one', 'metric_two'])
    self.assertIn(trace_summary.metric_bundles[1].specs[0].id,
                  ['metric_one', 'metric_two'])
    tp.close()

  def test_trace_summary_success_with_metadata_query(self):
    metric_spec = """metric_spec: {
        id: "memory_per_process"
        value: "dur"
        query: {
          simple_slices {
            process_name_glob: "ab*"
          }
        }
      }
      query: {
        id: "metadata_query"
        sql {
          sql: "SELECT \'foo\' AS key,  \'bar\' AS value"
          column_names: "key"
          column_names: "value"
        }
      }
      """
    tp = create_tp(trace=example_android_trace_path())
    trace_summary = tp.trace_summary([metric_spec], ['memory_per_process'],
                                     metadata_query_id='metadata_query')
    self.assertEqual(trace_summary.metric_bundles[0].specs[0].id,
                     'memory_per_process')
    self.assertTrue(hasattr(trace_summary, 'metadata'))
    tp.close()

  def test_trace_summary_dont_execute(self):
    metric_spec = """metric_spec: {
        id: "memory_per_process"
        value: "dur"
        query: {
          simple_slices {
            process_name_glob: "ab*"
          }
        }
      }
      """

    tp = create_tp(trace=example_android_trace_path())
    trace_summary = tp.trace_summary([metric_spec], [])
    self.assertEqual(len(trace_summary.metric_bundles), 0)
    tp.close()

  def test_trace_summary_no_ids_specified(self):
    metric_spec_1 = """metric_spec: {
        id: "metric_one"
        value: "dur"
        query: {
          simple_slices {
            process_name_glob: "ab*"
          }
        }
      }
      """
    metric_spec_2 = """metric_spec: {
        id: "metric_two"
        value: "ts"
        query: {
          simple_slices {
            process_name_glob: "cd*"
          }
        }
      }
      """
    tp = create_tp(trace=example_android_trace_path())
    trace_summary = tp.trace_summary([metric_spec_1, metric_spec_2])
    self.assertEqual(len(trace_summary.metric_bundles), 2)
    self.assertIn(trace_summary.metric_bundles[0].specs[0].id,
                  ['metric_one', 'metric_two'])
    self.assertIn(trace_summary.metric_bundles[1].specs[0].id,
                  ['metric_one', 'metric_two'])
    tp.close()

  def test_trace_summary_specs_as_bytes(self):
    platform_delegate = PLATFORM_DELEGATE()
    protos = ProtoFactory(platform_delegate)

    metric_spec_1 = protos.TraceSummarySpec()
    metric_1 = protos.TraceMetricV2Spec()
    metric_1.id = 'metric_one'
    metric_1.value = 'dur'
    metric_1.query.simple_slices.process_name_glob = 'ab*'
    metric_spec_1.metric_spec.extend([metric_1])
    metric_spec_1_bytes = metric_spec_1.SerializeToString()

    metric_spec_2 = protos.TraceSummarySpec()
    metric_2 = protos.TraceMetricV2Spec()
    metric_2.id = 'metric_two'
    metric_2.value = 'ts'
    metric_2.query.simple_slices.process_name_glob = 'cd*'
    metric_spec_2.metric_spec.extend([metric_2])
    metric_spec_2_bytes = metric_spec_2.SerializeToString()

    tp = create_tp(trace=example_android_trace_path())
    trace_summary = tp.trace_summary([metric_spec_1_bytes, metric_spec_2_bytes])
    self.assertEqual(len(trace_summary.metric_bundles), 2)
    self.assertIn(trace_summary.metric_bundles[0].specs[0].id,
                  ['metric_one', 'metric_two'])
    self.assertIn(trace_summary.metric_bundles[1].specs[0].id,
                  ['metric_one', 'metric_two'])
    tp.close()

  def test_trace_summary_specs_as_bytes_and_text(self):
    platform_delegate = PLATFORM_DELEGATE()
    protos = ProtoFactory(platform_delegate)

    metric_spec_1 = protos.TraceSummarySpec()
    metric_1 = protos.TraceMetricV2Spec()
    metric_1.id = 'metric_one'
    metric_1.value = 'dur'
    metric_1.query.simple_slices.process_name_glob = 'ab*'
    metric_spec_1.metric_spec.extend([metric_1])
    metric_spec_1_bytes = metric_spec_1.SerializeToString()

    metric_spec_2 = """metric_spec: {
        id: "metric_two"
        value: "ts"
        query: {
          simple_slices {
            process_name_glob: "cd*"
          }
        }
      }
      """

    tp = create_tp(trace=example_android_trace_path())
    trace_summary = tp.trace_summary([metric_spec_1_bytes, metric_spec_2])
    self.assertEqual(len(trace_summary.metric_bundles), 2)
    self.assertIn(trace_summary.metric_bundles[0].specs[0].id,
                  ['metric_one', 'metric_two'])
    self.assertIn(trace_summary.metric_bundles[1].specs[0].id,
                  ['metric_one', 'metric_two'])
    tp.close()

  def test_metadata_from_path(self):
    # When loading a trace directly from a path, metadata should be empty
    with create_tp(trace=example_android_trace_path()) as tp:
      self.assertEqual(tp.metadata, {"_path": example_android_trace_path()})

  def test_metadata_from_file(self):
    # When loading a trace from a file object, metadata should be empty
    with open(example_android_trace_path(), 'rb') as file:
      with create_tp(trace=file) as tp:
        self.assertEqual(tp.metadata, {})

  def test_metadata_from_generator(self):
    # When loading a trace from a generator, metadata should be empty
    def reader_generator():
      with open(example_android_trace_path(), 'rb') as file:
        yield file.read(1024)

    with create_tp(trace=reader_generator()) as tp:
      self.assertEqual(tp.metadata, {})

  def test_metadata_from_resolver(self):
    # Test that metadata is captured from a URI resolver
    registry = PLATFORM_DELEGATE().default_resolver_registry()
    registry.register(SimpleResolver)

    # Create a custom resolver that returns a single trace with known metadata
    class MetadataTestResolver(TraceUriResolver):
      PREFIX = 'metadata_test'

      def __init__(self):
        pass

      def resolve(self):
        return [
            TraceUriResolver.Result(
                example_android_trace_path(),
                metadata={
                    'test_key': 'test_value',
                    'trace_id': '12345'
                })
        ]

    registry.register(MetadataTestResolver)

    config = TraceProcessorConfig(
        bin_path=os.environ["SHELL_PATH"], resolver_registry=registry)

    with TraceProcessor(trace='metadata_test:', config=config) as tp:
      self.assertEqual(
          tp.metadata, {
              'test_key': 'test_value',
              'trace_id': '12345',
              '_path': example_android_trace_path()
          })

  def test_metadata_from_resolver_merged(self):
    # Test that metadata is merged when using nested resolvers
    registry = PLATFORM_DELEGATE().default_resolver_registry()

    # Create a two-level resolver to test metadata merging
    class OuterResolver(TraceUriResolver):
      PREFIX = 'outer'

      def __init__(self):
        pass

      def resolve(self):
        return [
            TraceUriResolver.Result(
                'inner:',
                metadata={
                    'outer_key': 'outer_value',
                    'shared_key': 'from_outer'
                })
        ]

    class InnerResolver(TraceUriResolver):
      PREFIX = 'inner'

      def __init__(self):
        pass

      def resolve(self):
        return [
            TraceUriResolver.Result(
                example_android_trace_path(),
                metadata={
                    'inner_key': 'inner_value',
                    'shared_key': 'from_inner'
                })
        ]

    registry.register(OuterResolver)
    registry.register(InnerResolver)

    config = TraceProcessorConfig(
        bin_path=os.environ["SHELL_PATH"], resolver_registry=registry)

    with TraceProcessor(trace='outer:', config=config) as tp:
      # Inner metadata should override outer metadata for shared keys
      expected_metadata = {
          'outer_key': 'outer_value',
          'inner_key': 'inner_value',
          'shared_key': 'from_inner',
          '_path': example_android_trace_path()
      }
      self.assertEqual(tp.metadata, expected_metadata)

  def test_export_perfetto(self):
    import json
    import tarfile
    with create_tp(trace=example_android_trace_path()) as tp:
      with tempfile.NamedTemporaryFile(suffix='.tar', delete=False) as f:
        output_path = f.name
      try:
        tp.export(output_path, 'perfetto')

        # Verify the output is a valid tar archive.
        self.assertTrue(tarfile.is_tarfile(output_path))

        with tarfile.open(output_path) as tf:
          names = tf.getnames()

          # The manifest must be the first member of the archive.
          self.assertEqual(names[0], 'perfetto_manifest.json')

          arrow_names = [n for n in names if n.endswith('.arrow')]
          self.assertGreater(len(arrow_names), 0)

          # Verify arrow files have the ARROW1 magic bytes.
          for name in arrow_names:
            member = tf.extractfile(name)
            self.assertIsNotNone(member)
            data = member.read()
            self.assertEqual(data[:6], b'ARROW1',
                             f'{name} missing ARROW1 header')

          # Verify the manifest contents match the archived arrow files.
          manifest_member = tf.extractfile('perfetto_manifest.json')
          self.assertIsNotNone(manifest_member)
          manifest = json.loads(manifest_member.read())['perfetto_manifest']
          self.assertEqual(manifest['version'], 1)

          table_files = [
              f for f in manifest['files'] if '__exported_table_schema' in f
          ]
          self.assertGreater(len(table_files), 0)
          for f in table_files:
            self.assertEqual(f['__exported_table_schema']['format'], 1)
          manifest_paths = {f['path'] for f in table_files}
          self.assertEqual(manifest_paths, set(arrow_names))
      finally:
        os.unlink(output_path)


class TestApiUnixSession(unittest.TestCase):
  """Tests TraceProcessor(remote=...) against a real `server unix` session."""

  SESSION_NAME = 'test-session'

  @classmethod
  def setUpClass(cls):
    # Point XDG_RUNTIME_DIR at a fresh temp dir so the session socket can't
    # collide with a real session on this machine. The server inherits this
    # environment, so both sides resolve the same socket path.
    tmp_dir = tempfile.TemporaryDirectory()
    cls.addClassCleanup(tmp_dir.cleanup)
    env_patch = mock.patch.dict(os.environ, {'XDG_RUNTIME_DIR': tmp_dir.name})
    env_patch.start()
    cls.addClassCleanup(env_patch.stop)

    stderr_file = tempfile.TemporaryFile()
    cls.addClassCleanup(stderr_file.close)
    cmd = [
        os.environ['SHELL_PATH'], 'server', 'unix', '--name', cls.SESSION_NAME,
        example_android_trace_path()
    ]
    server = subprocess.Popen(
        cmd, stdout=subprocess.PIPE, stderr=stderr_file, text=True)
    cls.addClassCleanup(cls._stop_server, server)

    # The server loads the trace, binds the socket and only then prints a
    # one-line startup record to stdout (its logs go to stderr). So any line
    # means the session is ready; EOF means the server exited.
    if not server.stdout.readline():
      # The server wrote through a shared file offset, so rewind before reading.
      stderr_file.seek(0)
      raise RuntimeError('trace_processor server unix failed to start:\n' +
                         stderr_file.read().decode())

  @staticmethod
  def _stop_server(server: subprocess.Popen):
    server.terminate()
    server.wait()
    server.stdout.close()

  def test_query(self):
    with TraceProcessor(remote=self.SESSION_NAME) as tp:
      qr_iterator = tp.query('select * from slice limit 10')
      dur_result = [
          178646, 119740, 58073, 155000, 173177, 20209377, 3589167, 90104,
          275312, 65313
      ]
      self.assertEqual([row.dur for row in qr_iterator], dur_result)

  def test_query_large_result(self):
    # Big enough that the server splits the result across many messages.
    num_rows = 100000
    with TraceProcessor(remote=self.SESSION_NAME) as tp:
      qr_iterator = tp.query(f'''
          with recursive n(i) as (
            select 0 union all select i + 1 from n where i < {num_rows - 1}
          )
          select i from n
      ''')
      self.assertEqual([row.i for row in qr_iterator], list(range(num_rows)))

  def test_query_error_then_recovers(self):
    with TraceProcessor(remote=self.SESSION_NAME) as tp:
      with self.assertRaisesRegex(TraceProcessorException, 'no_such_table'):
        tp.query('select * from no_such_table')
      # The error response must be fully consumed, leaving the connection
      # usable for the next query.
      self.assertEqual(next(tp.query('select 1 as x')).x, 1)

  def test_missing_session(self):
    with self.assertRaisesRegex(TraceProcessorException,
                                'No live trace processor session') as ctx:
      TraceProcessor(remote='no-such-session')
    self.assertIsInstance(ctx.exception.__cause__, FileNotFoundError)