#

# 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.

#

# pytype: skip-file

import base64

import logging

import unittest

import proto

import pytest

from google.protobuf import message

import apache_beam as beam

from apache_beam import typehints

from apache_beam.coders import coders

from apache_beam.coders import proto2_coder_test_messages_pb2 as test_message

from apache_beam.coders.avro_record import AvroRecord

from apache_beam.coders.typecoders import registry as coders_registry

from apache_beam.testing.test_pipeline import TestPipeline

class PickleCoderTest(unittest.TestCase):

def test_basics(self):

v = ('a' * 10, 'b' * 90)

pickler = coders.PickleCoder()

self.assertEqual(v, pickler.decode(pickler.encode(v)))

pickler = coders.Base64PickleCoder()

self.assertEqual(v, pickler.decode(pickler.encode(v)))

self.assertEqual(

coders.Base64PickleCoder().encode(v),

base64.b64encode(coders.PickleCoder().encode(v)))

def test_equality(self):

self.assertEqual(coders.PickleCoder(), coders.PickleCoder())

self.assertEqual(coders.Base64PickleCoder(), coders.Base64PickleCoder())

self.assertNotEqual(coders.Base64PickleCoder(), coders.PickleCoder())

self.assertNotEqual(coders.Base64PickleCoder(), object())

class CodersTest(unittest.TestCase):

def test_str_utf8_coder(self):

real_coder = coders_registry.get_coder(bytes)

expected_coder = coders.BytesCoder()

self.assertEqual(real_coder.encode(b'abc'), expected_coder.encode(b'abc'))

self.assertEqual(b'abc', real_coder.decode(real_coder.encode(b'abc')))

# The test proto message file was generated by running the following:

#

# `cd <beam repo>`

# `cp sdks/java/extensions/protobuf/src/test/proto/\

# proto2_coder_test_messages.proto sdks/python/apache_beam/coders/`

# `cd sdks/python`

# `protoc apache_beam/coders/proto2_coder_test_messages.proto

# --python_out=.

# `rm apache_beam/coders/proto2_coder_test_message.proto`

#

# Note: The protoc version should match the protobuf library version specified

# in setup.py.

#

# TODO(https://github.com/apache/beam/issues/22319): The proto file should be

# placed in a common directory that can be shared between java and python.

class ProtoCoderTest(unittest.TestCase):

def test_proto_coder(self):

ma = test_message.MessageA()

mb = ma.field2.add()

mb.field1 = True

ma.field1 = 'hello world'

expected_coder = coders.ProtoCoder(ma.__class__)

real_coder = coders_registry.get_coder(ma.__class__)

self.assertEqual(expected_coder, real_coder)

self.assertEqual(real_coder.encode(ma), expected_coder.encode(ma))

self.assertEqual(ma, real_coder.decode(real_coder.encode(ma)))

self.assertEqual(ma.__class__, real_coder.to_type_hint())

def test_proto_coder_on_protobuf_message_subclasses(self):

# This replicates a scenario where users provide message.Message as the

# output typehint for a Map function, even though the actual output messages

# are subclasses of message.Message.

ma = test_message.MessageA()

mb = ma.field2.add()

mb.field1 = True

ma.field1 = 'hello world'

coder = coders_registry.get_coder(message.Message)

# For messages of google.protobuf.message.Message, the fallback coder will

# be FastPrimitivesCoder rather than ProtoCoder.

# See the comment on ProtoCoder.from_type_hint() for further details.

self.assertEqual(coder, coders.FastPrimitivesCoder())

self.assertEqual(ma, coder.decode(coder.encode(ma)))

class DeterministicProtoCoderTest(unittest.TestCase):

def test_deterministic_proto_coder(self):

ma = test_message.MessageA()

mb = ma.field2.add()

mb.field1 = True

ma.field1 = 'hello world'

expected_coder = coders.DeterministicProtoCoder(ma.__class__)

real_coder = (

coders_registry.get_coder(

ma.__class__).as_deterministic_coder(step_label='unused'))

self.assertTrue(real_coder.is_deterministic())

self.assertEqual(expected_coder, real_coder)

self.assertEqual(real_coder.encode(ma), expected_coder.encode(ma))

self.assertEqual(ma, real_coder.decode(real_coder.encode(ma)))

def test_deterministic_proto_coder_determinism(self):

for _ in range(10):

keys = list(range(20))

mm_forward = test_message.MessageWithMap()

for key in keys:

mm_forward.field1[str(key)].field1 = str(key)

mm_reverse = test_message.MessageWithMap()

for key in reversed(keys):

mm_reverse.field1[str(key)].field1 = str(key)

coder = coders.DeterministicProtoCoder(mm_forward.__class__)

self.assertEqual(coder.encode(mm_forward), coder.encode(mm_reverse))

class ProtoPlusMessageB(proto.Message):

field1 = proto.Field(proto.BOOL, number=1)

class ProtoPlusMessageA(proto.Message):

field1 = proto.Field(proto.STRING, number=1)

field2 = proto.RepeatedField(ProtoPlusMessageB, number=2)

class ProtoPlusMessageWithMap(proto.Message):

field1 = proto.MapField(proto.STRING, ProtoPlusMessageA, number=1)

class ProtoPlusCoderTest(unittest.TestCase):

def test_proto_plus_coder(self):

ma = ProtoPlusMessageA()

ma.field2 = [ProtoPlusMessageB(field1=True)]

ma.field1 = 'hello world'

expected_coder = coders.ProtoPlusCoder(ma.__class__)

real_coder = coders_registry.get_coder(ma.__class__)

self.assertTrue(issubclass(ma.__class__, proto.Message))

self.assertEqual(expected_coder, real_coder)

self.assertTrue(real_coder.is_deterministic())

self.assertEqual(real_coder.encode(ma), expected_coder.encode(ma))

self.assertEqual(ma, real_coder.decode(real_coder.encode(ma)))

def test_proto_plus_coder_determinism(self):

for _ in range(10):

keys = list(range(20))

mm_forward = ProtoPlusMessageWithMap()

for key in keys:

mm_forward.field1[str(key)] = ProtoPlusMessageA(field1=str(key)) # pylint: disable=E1137

mm_reverse = ProtoPlusMessageWithMap()

for key in reversed(keys):

mm_reverse.field1[str(key)] = ProtoPlusMessageA(field1=str(key)) # pylint: disable=E1137

coder = coders.ProtoPlusCoder(ProtoPlusMessageWithMap)

self.assertEqual(coder.encode(mm_forward), coder.encode(mm_reverse))

class AvroTestCoder(coders.AvroGenericCoder):

SCHEMA = """

{

"type": "record", "name": "testrecord",

"fields": [

{"name": "name", "type": "string"},

{"name": "age", "type": "int"}

]

}

"""

def __init__(self):

super().__init__(self.SCHEMA)

class AvroTestRecord(AvroRecord):

pass

coders_registry.register_coder(AvroTestRecord, AvroTestCoder)

class AvroCoderTest(unittest.TestCase):

def test_avro_record_coder(self):

real_coder = coders_registry.get_coder(AvroTestRecord)

expected_coder = AvroTestCoder()

self.assertEqual(

real_coder.encode(

AvroTestRecord({

"name": "Daenerys targaryen", "age": 23

})),

expected_coder.encode(

AvroTestRecord({

"name": "Daenerys targaryen", "age": 23

})))

self.assertEqual(

AvroTestRecord({

"name": "Jon Snow", "age": 23

}),

real_coder.decode(

real_coder.encode(AvroTestRecord({

"name": "Jon Snow", "age": 23

}))))

class DummyClass(object):

"""A class with no registered coder."""

def __init__(self):

pass

def __eq__(self, other):

if isinstance(other, self.__class__):

return True

return False

def __hash__(self):

return hash(type(self))

class FallbackCoderTest(unittest.TestCase):

def test_default_fallback_path(self):

"""Test fallback path picks a matching coder if no coder is registered."""

coder = coders_registry.get_coder(DummyClass)

# No matching coder, so picks the last fallback coder which is a

# FastPrimitivesCoder.

self.assertEqual(coder, coders.FastPrimitivesCoder())

self.assertEqual(DummyClass(), coder.decode(coder.encode(DummyClass())))

class NullableCoderTest(unittest.TestCase):

def test_determinism(self):

deterministic = coders_registry.get_coder(typehints.Optional[int])

deterministic.as_deterministic_coder('label')

complex_deterministic = coders_registry.get_coder(

typehints.Optional[DummyClass])

complex_deterministic.as_deterministic_coder('label')

nondeterministic = coders.NullableCoder(coders.Base64PickleCoder())

with pytest.raises(ValueError):

nondeterministic.as_deterministic_coder('label')

class LengthPrefixCoderTest(unittest.TestCase):

def test_to_type_hint(self):

coder = coders.LengthPrefixCoder(coders.BytesCoder())

assert coder.to_type_hint() is bytes

class NumpyIntAsKeyTest(unittest.TestCase):

def test_numpy_int(self):

# this type is not supported as the key

import numpy as np

with self.assertRaisesRegex(Exception, "Unable to deterministically"):

with TestPipeline() as p:

indata = p | "Create" >> beam.Create([(a, int(a))

for a in np.arange(3)])

# Apply CombinePerkey to sum values for each key.

_ = indata | "CombinePerKey" >> beam.CombinePerKey(sum)

class WindowedValueCoderTest(unittest.TestCase):

def test_to_type_hint(self):

coder = coders.WindowedValueCoder(coders.VarIntCoder())

self.assertEqual(coder.to_type_hint(), typehints.WindowedValue[int]) # type: ignore[misc]

if __name__ == '__main__':

logging.getLogger().setLevel(logging.INFO)

unittest.main()