sdks/python/apache_beam/typehints/sharded_key_type_test.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#

"""Unit tests for the ShardedKeyTypeConstraint."""

# pytype: skip-file

from __future__ import absolute_import

from apache_beam.typehints import Tuple
from apache_beam.typehints import typehints
from apache_beam.typehints.sharded_key_type import ShardedKeyType
from apache_beam.typehints.typehints_test import TypeHintTestCase
from apache_beam.utils.sharded_key import ShardedKey


class ShardedKeyTypeConstraintTest(TypeHintTestCase):
  def test_compatibility(self):
    constraint1 = ShardedKeyType[int]
    constraint2 = ShardedKeyType[str]

    self.assertCompatible(constraint1, constraint1)
    self.assertCompatible(constraint2, constraint2)
    self.assertNotCompatible(constraint1, constraint2)

  def test_repr(self):
    constraint = ShardedKeyType[int]
    self.assertEqual('ShardedKey[int]', repr(constraint))

  def test_type_check_not_sharded_key(self):
    constraint = ShardedKeyType[int]
    obj = 5
    with self.assertRaises(TypeError) as e:
      constraint.type_check(obj)
    self.assertEqual(
        "ShardedKey type-constraint violated. Valid object instance must be of "
        "type 'ShardedKey'. Instead, an instance of 'int' was received.",
        e.exception.args[0])

  def test_type_check_invalid_key_type(self):
    constraint = ShardedKeyType[int]
    obj = ShardedKey(key='abc', shard_id=b'123')
    with self.assertRaises((TypeError, TypeError)) as e:
      constraint.type_check(obj)
    self.assertEqual(
        "ShardedKey[int] type-constraint violated. The type of key in "
        "'ShardedKey' is incorrect. Expected an instance of type 'int', "
        "instead received an instance of type 'str'.",
        e.exception.args[0])

  def test_type_check_valid_simple_type(self):
    constraint = ShardedKeyType[str]
    obj = ShardedKey(key='abc', shard_id=b'123')
    self.assertIsNone(constraint.type_check(obj))

  def test_type_check_valid_composite_type(self):
    constraint = ShardedKeyType[Tuple[int, str]]
    obj = ShardedKey(key=(1, 'a'), shard_id=b'123')
    self.assertIsNone(constraint.type_check(obj))

  def test_match_type_variables(self):
    K = typehints.TypeVariable('K')  # pylint: disable=invalid-name
    constraint = ShardedKeyType[K]
    self.assertEqual({K: int},
                     constraint.match_type_variables(ShardedKeyType[int]))

  def test_getitem(self):
    K = typehints.TypeVariable('K')  # pylint: disable=invalid-name
    T = typehints.TypeVariable('T')  # pylint: disable=invalid-name
    with self.assertRaisesRegex(TypeError,
                                'Parameter to ShardedKeyType hint.*'):
      _ = ShardedKeyType[K, T]
    with self.assertRaisesRegex(TypeError,
                                'Parameter to ShardedKeyType hint.*'):
      _ = ShardedKeyType[(K, T)]