src/dctv/test_cmd_query.py

# 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.
"""Tests for command query module"""
import logging
import json
import csv
from os.path import join as pjoin
from tempfile import NamedTemporaryFile

import pytest

from .cmd_query import run_query
from .util import dctv_dir

log = logging.getLogger(__name__)

# pylint: disable=missing-docstring,line-too-long

def _do_query(query, output_format, block_size, **mounts):
  with NamedTemporaryFile(mode="w+",
                          prefix="dctv-cmd-query-test",
                          suffix=".json") as ntf:
    run_query(
      query=query,
      trace_mounts=[
        (mount_path.split("."), file_name)
        for mount_path, file_name in mounts.items()],
      out_file=ntf,
      output_format=output_format,
      block_size=block_size,
      temp_hack_lenient_metadata=True)
    ntf.seek(0)
    if output_format != "csv_rows":
      return json.loads(ntf.read())
    dialect = csv.Sniffer().sniff(ntf.read(1024))
    ntf.seek(0)
    reader = csv.reader(ntf, dialect)
    return list(reader)

@pytest.mark.parametrize("block_size", [1, 2, None])
def test_cmd_query(block_size):
  # pylint: disable=bad-whitespace
  assert _do_query("SELECT 4",
                   "json_columns", block_size) \
    == {"4": [4]}
  assert _do_query("VALUES (1), (2), (3)",
                   "json_columns", block_size) == \
    {"col0": [1, 2, 3]}
  assert _do_query("VALUES ('foo', 1), ('bar', 2), ('qux', 3)",
                   "json_columns", block_size) == \
    {
      "col0": ["foo", "bar", "qux"],
      "col1": [1, 2, 3],
    }
  q = "(VALUES (1.1, 'qwer'), (2, 'foo'), (3, NULL))"
  assert _do_query(q, "json_columns", block_size) == {
    "col0": [1.1, 2, 3],
    "col1": ["qwer", "foo", None],
  }
  assert _do_query(q, "json_records", block_size) == [
    {"col0": 1.1, "col1": "qwer"},
    {"col0": 2,   "col1": "foo"},
    {"col0": 3,   "col1": None},
  ]
  assert _do_query(q, "csv_rows", block_size) == [
    ["col0", "col1"],
    ["1.1", "qwer"],
    ["2.0", "foo"],
    ["3.0", ""],
  ]


def test_cmd_query_file():
  test_trace_file_name = pjoin(dctv_dir(),
                               "test-data",
                               "dragonball.mini.trace")
  assert _do_query("SELECT COUNT(*) AS c FROM foo.raw_events.sched_switch",
                   "json_columns", None,
                   foo=test_trace_file_name) == {"c": [32370]}