src/dctv/test_unop.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.
"""Test cases for SQL unary operations"""

# pylint: disable=missing-docstring

import logging
import pytest

from .test_binop import (
  get_example_tbl,
)

from .query import (
  InvalidQueryException,
  _UNOP,
)

from .test_query import (
  ALL_SCHEMAS,
  _execute_qt,
)

from .sql import (
  UnaryOperation,
)

from .test_sql import (
  ps2qt,
)

log = logging.getLogger(__name__)

_TEST_UNOP = ("+ - ~ not ! ceiling() "
              "floor() round() trunc()").split()

@pytest.mark.parametrize("operator_name", _TEST_UNOP)
@pytest.mark.parametrize("right", ["2", "E", "N"])
@pytest.mark.parametrize("schema_name", ALL_SCHEMAS)
def test_expression_unop(operator_name, right, schema_name):
  tbl = get_example_tbl(schema_name, right)
  schema = tbl["vl"].schema
  if operator_name.endswith("()"):
    operator_name = operator_name[:-2]
    q = "SELECT {}(vl) AS x FROM tbl" .format(operator_name)
  else:
    q = "SELECT {} vl AS x FROM tbl" .format(operator_name)
  munged_operator_name = UnaryOperation.EQUIV.get(
    operator_name, operator_name)
  operator = _UNOP[munged_operator_name]

  # pylint: disable=redefined-variable-type

  should_throw = False
  if (not should_throw and
      schema.is_string and not operator.supports_strings):
    should_throw = "string"
  if (not should_throw and
      schema.dtype.kind not in operator.supported_kinds):
    should_throw = "kind"
  if (not should_throw and
      schema.domain and not operator.supports_domains):
    should_throw = "domain"

  try:
    qt = ps2qt(q, tbl=tbl)
    _execute_qt(qt, as_raw_array=True)["x"]
  except InvalidQueryException:
    if not should_throw:
      raise
    return
  assert not should_throw