#!/usr/bin/env python
"""Test RDFStruct implementations."""


import base64
import random
from typing import Text

from absl import app
from absl.testing import absltest

from google.protobuf import any_pb2
from google.protobuf import wrappers_pb2
from grr_response_core.lib import rdfvalue
from grr_response_core.lib import type_info
from grr_response_core.lib.rdfvalues import client as rdf_client
from grr_response_core.lib.rdfvalues import client_fs as rdf_client_fs
from grr_response_core.lib.rdfvalues import client_network as rdf_client_network
from grr_response_core.lib.rdfvalues import client_stats as rdf_client_stats
from grr_response_core.lib.rdfvalues import flows as rdf_flows
from grr_response_core.lib.rdfvalues import paths as rdf_paths
from grr_response_core.lib.rdfvalues import structs as rdf_structs
from grr_response_core.lib.rdfvalues import test_base as rdf_test_base
from grr_response_proto import tests_pb2
from grr_response_server.rdfvalues import flow_objects as rdf_flow_objects
from grr.test_lib import test_lib

# pylint: mode=test


class TestStructWithManyFields(rdf_structs.RDFProtoStruct):
  """A test struct object."""

  type_description = type_info.TypeDescriptorSet(*[
      rdf_structs.ProtoString(
          name="foobar_%d" % i,
          field_number=i + 1,
          default="string",
          description="A string value",
      ) for i in range(100)
  ])


class TestStruct(rdf_structs.RDFProtoStruct):
  """A test struct object."""

  type_description = type_info.TypeDescriptorSet(
      rdf_structs.ProtoString(
          name="foobar",
          field_number=1,
          default="string",
          description="A string value",
          labels=[rdf_structs.SemanticDescriptor.Labels.HIDDEN]),
      rdf_structs.ProtoUnsignedInteger(
          name="int", field_number=2, default=5,
          description="An integer value"),
      rdf_structs.ProtoList(
          rdf_structs.ProtoString(
              name="repeated",
              field_number=3,
              description="A repeated string value")),

      # We can serialize an arbitrary RDFValue. This will be serialized into a
      # binary string and parsed on demand.
      rdf_structs.ProtoRDFValue(
          name="urn",
          field_number=6,
          default=rdfvalue.RDFURN("www.google.com"),
          rdf_type="RDFURN",
          description="An arbitrary RDFValue field."),
      rdf_structs.ProtoEnum(
          name="type",
          field_number=7,
          enum_name="Type",
          enum={
              "FIRST": 1,
              "SECOND": 2,
              "THIRD": 3
          },
          default=3,
          description="An enum field"),
      rdf_structs.ProtoFloat(
          name="float",
          field_number=8,
          description="A float number",
          default=1.1),
  )


# In order to define a recursive structure we must add it manually after the
# class definition.
TestStruct.AddDescriptor(
    rdf_structs.ProtoEmbedded(name="nested", field_number=4,
                              nested=TestStruct),)

TestStruct.AddDescriptor(
    rdf_structs.ProtoList(
        rdf_structs.ProtoEmbedded(
            name="repeat_nested", field_number=5, nested=TestStruct)),)


class VersionedTestStructV1(rdf_structs.RDFProtoStruct):
  """A test struct object."""

  type_description = type_info.TypeDescriptorSet(
      rdf_structs.ProtoBoolean(
          name="bool1",
          field_number=1,
          default=False,
      ))


class VersionedTestStructV2(rdf_structs.RDFProtoStruct):
  """A test struct object."""

  type_description = type_info.TypeDescriptorSet(
      rdf_structs.ProtoBoolean(
          name="bool1",
          field_number=1,
          default=False,
      ), rdf_structs.ProtoBoolean(
          name="bool2",
          field_number=2,
          default=False,
      ))


class TestStructWithBool(rdf_structs.RDFProtoStruct):
  """A test struct object."""

  type_description = type_info.TypeDescriptorSet(
      rdf_structs.ProtoBoolean(
          name="foo",
          field_number=1,
          default=False,
      ))


class TestStructWithEnum(rdf_structs.RDFProtoStruct):
  """A test struct object."""

  type_description = type_info.TypeDescriptorSet(
      rdf_structs.ProtoEnum(
          name="foo",
          field_number=1,
          enum_name="Foo",
          enum={
              "ZERO": 0,
              "ONE": 1,
              "TWO": 2
          },
          default=0))


class PartialTest1(rdf_structs.RDFProtoStruct):
  """This is a protobuf with fewer fields than TestStruct."""
  type_description = type_info.TypeDescriptorSet(
      rdf_structs.ProtoUnsignedInteger(name="int", field_number=2),)


class DynamicTypeTest(rdf_structs.RDFProtoStruct):
  """A protobuf with dynamic types."""

  type_description = type_info.TypeDescriptorSet(
      rdf_structs.ProtoString(
          name="type",
          field_number=1,
          # By default return the TestStruct proto.
          default="TestStruct",
          description="A string value"),
      rdf_structs.ProtoDynamicEmbedded(
          name="dynamic",
          # The callback here returns the type specified by the type member.
          dynamic_cb=lambda x: rdf_structs.RDFProtoStruct.classes.get(x.type),
          field_number=2,
          description="A dynamic value based on another field."),
      rdf_structs.ProtoEmbedded(
          name="nested", field_number=3, nested=rdf_client.User))


class DynamicAnyValueTypeTest(rdf_structs.RDFProtoStruct):
  """A protobuf with dynamic types stored in AnyValue messages."""

  type_description = type_info.TypeDescriptorSet(
      rdf_structs.ProtoString(
          name="type", field_number=1, description="A string value"),
      rdf_structs.ProtoDynamicAnyValueEmbedded(
          name="dynamic",
          # The callback here returns the type specified by the type member.
          dynamic_cb=lambda x: rdf_structs.RDFProtoStruct.classes.get(x.type),
          field_number=2,
          description="A dynamic value based on another field."),
  )


class AnyValueWithoutTypeFunctionTest(rdf_structs.RDFProtoStruct):
  """A protobuf with dynamic types stored in AnyValue messages."""

  type_description = type_info.TypeDescriptorSet(
      rdf_structs.ProtoDynamicAnyValueEmbedded(
          name="dynamic", field_number=1, description="A dynamic value."),)


class LateBindingTest(rdf_structs.RDFProtoStruct):
  type_description = type_info.TypeDescriptorSet(
      # A nested protobuf referring to an undefined type.
      rdf_structs.ProtoEmbedded(
          name="nested", field_number=1, nested="UndefinedYet"),
      rdf_structs.ProtoRDFValue(
          name="rdfvalue",
          field_number=6,
          rdf_type="UndefinedRDFValue",
          description="An undefined RDFValue field."),

      # A repeated late bound field.
      rdf_structs.ProtoList(
          rdf_structs.ProtoRDFValue(
              name="repeated",
              field_number=7,
              rdf_type="UndefinedRDFValue2",
              description="An undefined RDFValue field.")),
  )


class UnionTest(rdf_structs.RDFProtoStruct):
  union_field = "struct_flavor"

  type_description = type_info.TypeDescriptorSet(
      rdf_structs.ProtoEnum(
          name="struct_flavor",
          field_number=1,
          enum_name="Type",
          enum={
              "FIRST": 1,
              "SECOND": 2,
              "THIRD": 3
          },
          default=3,
          description="An union enum field"),
      rdf_structs.ProtoFloat(
          name="first",
          field_number=2,
          description="A float number",
          default=1.1),
      rdf_structs.ProtoString(
          name="second",
          field_number=3,
          default="string",
          description="A string value"),
      rdf_structs.ProtoUnsignedInteger(
          name="third",
          field_number=4,
          default=5,
          description="An integer value"),
  )


class RDFStructsTest(rdf_test_base.RDFValueTestMixin, test_lib.GRRBaseTest):
  """Test the RDFStruct implementation."""

  rdfvalue_class = TestStruct

  def CheckRDFValue(self, value, sample):
    # Verify that floats are almost equal, since serialization and conversion
    # of floats is not stable.
    self.assertAlmostEqual(value.float, sample.float)

    # Verify that all other fields are exactly equal.
    value_copy = value.Copy()
    value_copy.float = None

    sample_copy = sample.Copy()
    sample_copy.float = None

    super().CheckRDFValue(value_copy, sample_copy)

  def GenerateSample(self, number=1):
    return self.rdfvalue_class(
        int=number,
        foobar="foo%s" % number,
        urn="www.example.com",
        float=2.3 + number)

  def testDynamicType(self):
    test_pb = DynamicTypeTest()
    # We can not assign arbitrary values to the dynamic field.
    self.assertRaises(ValueError, setattr, test_pb, "dynamic", "hello")

    # Can assign a nested field.
    test_pb.dynamic.foobar = "Hello"
    self.assertIsInstance(test_pb.dynamic, TestStruct)

    # Test serialization/deserialization.
    serialized = test_pb.SerializeToBytes()
    self.assertEqual(DynamicTypeTest.FromSerializedBytes(serialized), test_pb)
    test_struct_class_repr = "'TestStruct'"

    expected_emitted_proto = """\
message DynamicTypeTest {{

  // A string value
  optional string type = 1 [default = {test_struct}];

  // A dynamic value based on another field.
  optional bytes dynamic = 2;
  optional User nested = 3;
}}
""".format(test_struct=test_struct_class_repr)

    # Test proto definition.
    self.assertEqual(DynamicTypeTest.EmitProto(), expected_emitted_proto)

  def testAnyValueWithoutTypeCallback(self):
    test_pb = AnyValueWithoutTypeFunctionTest()

    for value_to_assign in [
        rdfvalue.RDFString("test"),
        rdfvalue.RDFInteger(1234),
        rdfvalue.RDFBytes(b"abc"),
        rdf_flows.GrrStatus(status="WORKER_STUCK", error_message="stuck")
    ]:
      test_pb.dynamic = value_to_assign
      serialized = test_pb.SerializeToBytes()
      deserialized = AnyValueWithoutTypeFunctionTest.FromSerializedBytes(
          serialized)
      self.assertEqual(deserialized, test_pb)
      self.assertEqual(type(deserialized.dynamic), type(value_to_assign))

  def testDynamicAnyValueType(self):
    test_pb = DynamicAnyValueTypeTest(type="TestStruct")
    # We can not assign arbitrary values to the dynamic field.
    self.assertRaises(ValueError, setattr, test_pb, "dynamic", "hello")

    # Can assign a nested field.
    test_pb.dynamic.foobar = "Hello"
    self.assertIsInstance(test_pb.dynamic, TestStruct)

    # Test serialization/deserialization.
    serialized = test_pb.SerializeToBytes()
    self.assertEqual(
        DynamicAnyValueTypeTest.FromSerializedBytes(serialized), test_pb)

    # Test proto definition.
    self.assertEqual(
        DynamicAnyValueTypeTest.EmitProto(),
        "message DynamicAnyValueTypeTest {\n\n  "
        "// A string value\n  optional string type = 1;\n\n  "
        "// A dynamic value based on another field.\n  "
        "optional google.protobuf.Any dynamic = 2;\n}\n")

  def testUninitializedDynamicValueIsNonePerDefault(self):
    response = rdf_flow_objects.FlowResponse()  # Do not set payload.
    self.assertIsNone(response.payload)

  def testDynamicAnyValueTypeWithPrimitiveValues(self):
    test_pb = DynamicAnyValueTypeTest(type="RDFString")
    # We can not assign arbitrary values to the dynamic field.
    self.assertRaises(ValueError, setattr, test_pb, "dynamic", 42)

    # Can assign a nested field.
    test_pb.dynamic = rdfvalue.RDFString("Hello")
    self.assertIsInstance(test_pb.dynamic, rdfvalue.RDFString)

    # Test serialization/deserialization.
    serialized = test_pb.SerializeToBytes()
    unserialized = DynamicAnyValueTypeTest.FromSerializedBytes(serialized)
    self.assertEqual(unserialized, test_pb)
    self.assertEqual(unserialized.dynamic, "Hello")

  def testStructDefinition(self):
    """Ensure that errors in struct definitions are raised."""
    # A descriptor without a field number should raise.
    self.assertRaises(
        type_info.TypeValueError, rdf_structs.ProtoEmbedded, name="name")

    # Adding a duplicate field number should raise.
    self.assertRaises(
        type_info.TypeValueError, TestStruct.AddDescriptor,
        rdf_structs.ProtoUnsignedInteger(name="int", field_number=2))

    # Adding a descriptor which is not a Proto* descriptor is not allowed for
    # Struct fields:
    self.assertRaises(type_info.TypeValueError, TestStruct.AddDescriptor,
                      type_info.String(name="int"))

  def testRepeatedMember(self):
    tested = TestStruct(int=5)
    tested.foobar = "Hello"
    tested.repeated.Append("Good")
    tested.repeated.Append("Bye")

    for i in range(10):
      tested.repeat_nested.Append(foobar="Nest%s" % i)

    data = tested.SerializeToBytes()

    # Parse it again.
    new_tested = TestStruct.FromSerializedBytes(data)

    # Test the repeated field.
    self.assertLen(new_tested.repeat_nested, 10)
    self.assertEqual(new_tested.repeat_nested[1].foobar, "Nest1")

    # Check that slicing works.
    sliced = new_tested.repeat_nested[3:5]
    self.assertEqual(sliced.__class__, new_tested.repeat_nested.__class__)
    self.assertEqual(sliced.type_descriptor,
                     new_tested.repeat_nested.type_descriptor)

    self.assertLen(sliced, 2)
    self.assertEqual(sliced[0].foobar, "Nest3")

  def testRepeatedFieldHelperEqualToSequences(self):
    struct = TestStruct(repeated=["a", "b", "c"])

    self.assertNotEqual(struct.repeated, "abc")
    self.assertNotEqual(struct.repeated, {"a", "b", "c"})
    self.assertNotEqual(struct.repeated, {"a": "a", "b": "b", "c": "c"})

    self.assertNotEqual(struct.repeated, ["a"])
    self.assertNotEqual(struct.repeated, ["a", "b"])
    self.assertNotEqual(struct.repeated, ["a", "c", "b"])
    self.assertEqual(struct.repeated, ["a", "b", "c"])

    self.assertNotEqual(struct.repeated, ("a",))
    self.assertNotEqual(struct.repeated, ("a", "b"))
    self.assertNotEqual(struct.repeated, ("a", "c", "b"))
    self.assertEqual(struct.repeated, ("a", "b", "c"))

    self.assertIs(struct.repeated.__eq__(dict()), NotImplemented)
    self.assertIs(struct.repeated.__eq__(set()), NotImplemented)
    self.assertIs(struct.repeated.__eq__(b""), NotImplemented)
    self.assertIs(struct.repeated.__eq__(""), NotImplemented)
    self.assertIs(struct.repeated.__eq__(object()), NotImplemented)

  def testUnknownFields(self):
    """Test that unknown fields are preserved across decode/encode cycle."""
    tested = TestStruct(foobar="hello", int=5)
    tested.nested.foobar = "goodbye"

    self.assertEqual(tested.foobar, "hello")
    self.assertEqual(tested.nested.foobar, "goodbye")

    data = tested.SerializeToBytes()

    # Now unpack using a protobuf with less fields defined.
    reduced_tested = PartialTest1.FromSerializedBytes(data)

    self.assertEqual(reduced_tested.int, 5)
    self.assertRaises(AttributeError, getattr, reduced_tested, "foobar")

    # Re-Serialize using the simple protobuf.
    data2 = reduced_tested.SerializeToBytes()
    decoded_tested = TestStruct.FromSerializedBytes(data2)

    # The foobar field should have been preserved, despite PartialTest1() not
    # understanding it.
    self.assertEqual(decoded_tested.foobar, "hello")

    # Check that nested fields are also preserved.
    self.assertEqual(decoded_tested.nested.foobar, "goodbye")

  def testRDFStruct(self):
    tested = TestStruct()

    # can't set integers for string attributes.
    self.assertRaises(type_info.TypeValueError, setattr, tested, "foobar", 1)

    # This is a string so a string assignment is good:
    tested.foobar = "Hello"
    self.assertEqual(tested.foobar, "Hello")

    # This field must be another TestStruct instance..
    self.assertRaises(ValueError, setattr, tested, "nested", "foo")

    # Its ok to assign a compatible semantic protobuf.
    tested.nested = TestStruct(foobar="nested_foo")

    # Not OK to use the wrong semantic type.
    self.assertRaises(ValueError, setattr, tested, "nested",
                      PartialTest1(int=1))

    # Not OK to assign a serialized string - even if it is for the right type -
    # since there is no type checking.
    serialized = TestStruct(foobar="nested_foo").SerializeToBytes()
    self.assertRaises(ValueError, setattr, tested, "nested", serialized)

    # Nested accessors.
    self.assertEqual(tested.nested.foobar, "nested_foo")

    # Test repeated elements:

    # Empty list is ok:
    tested.repeated = []
    self.assertEqual(tested.repeated, [])

    tested.repeated = ["string"]
    self.assertEqual(tested.repeated, ["string"])

    self.assertRaises(type_info.TypeValueError, setattr, tested, "repeated",
                      [1, 2, 3])

    # Coercing on assignment. This field is an RDFURN:
    tested.urn = "www.example.com"
    self.assertIsInstance(tested.urn, rdfvalue.RDFURN)

    self.assertEqual(tested.urn, rdfvalue.RDFURN("www.example.com"))

    # Test enums.
    self.assertEqual(tested.type, 3)
    self.assertEqual(tested.type.name, "THIRD")

    tested.type = "FIRST"
    self.assertEqual(tested.type, 1)

    # Check that string assignments are case-insensitive.
    tested.type = "second"
    self.assertEqual(tested.type, 2)
    tested.type = "ThIrD"
    self.assertEqual(tested.type, 3)

    # Non-valid types are rejected.
    self.assertRaises(type_info.TypeValueError, setattr, tested, "type", "Foo")

    # Strings of digits should be accepted.
    tested.type = "2"
    self.assertEqual(tested.type, 2)
    # unicode strings should be treated the same way.
    tested.type = u"2"
    self.assertEqual(tested.type, 2)
    # Out of range values are permitted and preserved through serialization.
    tested.type = 4
    self.assertEqual(tested.type, 4)
    serialized_type = str(tested.type).encode("utf-8")
    tested.type = 1
    tested.type = serialized_type
    self.assertEqual(tested.type, 4)

  def testCacheInvalidation(self):
    path = rdf_paths.PathSpec(path="/", pathtype=rdf_paths.PathSpec.PathType.OS)
    for x in "01234":
      path.last.Append(path=x, pathtype=rdf_paths.PathSpec.PathType.OS)

    serialized = path.SerializeToBytes()

    path = rdf_paths.PathSpec.FromSerializedBytes(serialized)

    # At this point the wire format cache is fully populated (since the proto
    # had been parsed). We change a deeply nested member.
    path.last.path = "booo"

    # When we serialize the modified proto we should get the new field
    # serialized. If the cache is not properly invalidated, we will return the
    # old result instead.
    self.assertIn(b"booo", path.SerializeToBytes())

  def testLateBinding(self):
    # The LateBindingTest protobuf is not fully defined.
    self.assertRaises(KeyError, LateBindingTest.type_infos.__getitem__,
                      "nested")

    self.assertIn("UndefinedYet", rdfvalue._LATE_BINDING_STORE)

    # We can still use this protobuf
    tested = LateBindingTest()

    # But it does not know about the field yet.
    self.assertRaises(AttributeError, tested.Get, "nested")
    self.assertRaises(AttributeError, LateBindingTest, nested=None)

    # Now define the class. This should resolve the late bound fields and re-add
    # them to their owner protobufs.
    class UndefinedYet(rdf_structs.RDFProtoStruct):
      type_description = type_info.TypeDescriptorSet(
          rdf_structs.ProtoString(
              name="foobar", field_number=1, description="A string value"),)

    # The field is now resolved.
    self.assertNotIn("UndefinedYet", rdfvalue._LATE_BINDING_STORE)
    nested_field = LateBindingTest.type_infos["nested"]
    self.assertEqual(nested_field.name, "nested")

    # We can now use the protobuf as normal.
    tested = LateBindingTest()
    tested.nested.foobar = "foobar string"
    self.assertIsInstance(tested.nested, UndefinedYet)

  def testRDFValueLateBinding(self):
    # The LateBindingTest protobuf is not fully defined.
    self.assertRaises(KeyError, LateBindingTest.type_infos.__getitem__,
                      "rdfvalue")

    self.assertIn("UndefinedRDFValue", rdfvalue._LATE_BINDING_STORE)

    # We can still use this protobuf
    tested = LateBindingTest()

    # But it does not know about the field yet.
    self.assertRaises(AttributeError, tested.Get, "rdfvalue")
    self.assertRaises(AttributeError, LateBindingTest, rdfvalue="foo")

    # Now define the class. This should resolve the late bound fields and re-add
    # them to their owner protobufs.
    class UndefinedRDFValue(rdfvalue.RDFString):
      pass

    # The field is now resolved.
    self.assertNotIn("UndefinedRDFValue", rdfvalue._LATE_BINDING_STORE)
    rdfvalue_field = LateBindingTest.type_infos["rdfvalue"]
    self.assertEqual(rdfvalue_field.name, "rdfvalue")

    # We can now use the protobuf as normal.
    tested = LateBindingTest(rdfvalue="foo")
    self.assertEqual(type(tested.rdfvalue), UndefinedRDFValue)

  def testRepeatedRDFValueLateBinding(self):
    # The LateBindingTest protobuf is not fully defined.
    self.assertRaises(KeyError, LateBindingTest.type_infos.__getitem__,
                      "repeated")

    self.assertIn("UndefinedRDFValue2", rdfvalue._LATE_BINDING_STORE)

    # We can still use this protobuf
    tested = LateBindingTest()

    # But it does not know about the field yet.
    self.assertRaises(AttributeError, tested.Get, "repeated")
    self.assertRaises(AttributeError, LateBindingTest, repeated=["foo"])

    # Now define the class. This should resolve the late bound fields and re-add
    # them to their owner protobufs.
    class UndefinedRDFValue2(rdfvalue.RDFString):
      pass

    # The field is now resolved.
    self.assertNotIn("UndefinedRDFValue2", rdfvalue._LATE_BINDING_STORE)
    rdfvalue_field = LateBindingTest.type_infos["repeated"]
    self.assertEqual(rdfvalue_field.name, "repeated")

    # We can now use the protobuf as normal.
    tested = LateBindingTest(repeated=["foo"])
    self.assertLen(tested.repeated, 1)
    self.assertEqual(tested.repeated[0], "foo")
    self.assertEqual(type(tested.repeated[0]), UndefinedRDFValue2)

  def testRDFValueParsing(self):
    stat = rdf_client_fs.StatEntry.protobuf(st_mode=16877)
    data = stat.SerializeToString()

    result = rdf_client_fs.StatEntry.FromSerializedBytes(data)

    self.assertIsInstance(result.st_mode, rdf_client_fs.StatMode)

  def testDefaults(self):
    """Accessing a field which does not exist returns a default."""
    # An empty protobuf.
    tested = TestStruct()

    # Simple strings.
    self.assertFalse(tested.HasField("foobar"))
    self.assertEqual(tested.foobar, "string")
    self.assertFalse(tested.HasField("foobar"))

    # RDFValues.
    self.assertFalse(tested.HasField("urn"))
    self.assertEqual(tested.urn, "www.google.com")
    self.assertFalse(tested.HasField("urn"))

    # Nested fields: Accessing a nested field will create the nested protobuf.
    self.assertFalse(tested.HasField("nested"))
    self.assertEqual(tested.nested.urn, "www.google.com")
    self.assertTrue(tested.HasField("nested"))
    self.assertFalse(tested.nested.HasField("urn"))

  def testUnionCast(self):
    """Check union structs handling."""
    # An empty protobuf.
    tested_non_union = TestStruct()

    # Raises if union-casting is attempted on non-union proto.
    self.assertRaises(AttributeError, tested_non_union.UnionCast)

    # A proto with a semantic union_field. In this particular proto the chosen
    # union variant is called struct_flavor, but it's arbitrary, we also use
    # eg rule_type referring to the chosen union variant elsewhere. This is
    # determined by the value of union_field.
    tested_union = UnionTest()

    # Returns the default value of the default flavor if both struct_flavor and
    # the flavored field are not set.
    self.assertEqual(tested_union.UnionCast(), 5)

    tested_union.struct_flavor = "FIRST"
    # Returns the default value of the selected flavor if it's not been changed.
    self.assertEqual(tested_union.UnionCast(), 1.1)

    tested_union.struct_flavor = "SECOND"
    # Same as above.
    self.assertEqual(tested_union.UnionCast(), "string")

    tested_union.struct_flavor = "THIRD"
    # Once again.
    self.assertEqual(tested_union.UnionCast(), 5)

    tested_union.third = 1337
    # Returns the set flavor value.
    self.assertEqual(tested_union.UnionCast(), 1337)

    tested_union.first = 1.61803399
    # Raises if there is a flavored field set that doesn't match struct_flavor.
    self.assertRaises(ValueError, tested_union.UnionCast)

  def testClearFieldsWithLabelWorksCorrectly(self):
    t = TestStruct(foobar="foo", int=42)
    t.ClearFieldsWithLabel(rdf_structs.SemanticDescriptor.Labels.HIDDEN)
    self.assertFalse(t.HasField("foobar"))
    self.assertTrue(t.HasField("int"))

  def testClearFieldsWithLabelWorksCorrectlyOnNestedStructures(self):
    t = TestStruct(foobar="foo", int=42)
    t.nested = TestStruct(foobar="bar", int=43)
    t.ClearFieldsWithLabel(rdf_structs.SemanticDescriptor.Labels.HIDDEN)
    self.assertFalse(t.nested.HasField("foobar"))
    self.assertTrue(t.nested.HasField("int"))

  def testConversionToPrimitiveDictNoSerialization(self):
    test_struct = TestStruct(
        foobar="foo",
        int=2,
        repeated=["value0", "value1"],
        nested=TestStruct(int=567),
        repeat_nested=[TestStruct(int=568)])
    expected_dict = {
        "foobar": "foo",
        "int": 2,
        "repeated": ["value0", "value1"],
        "nested": {
            "int": 567
        },
        "repeat_nested": [{
            "int": 568
        }]
    }
    self.assertEqual(test_struct.ToPrimitiveDict(), expected_dict)

  def testConversionToPrimitiveDictWithSerialization(self):
    test_struct = TestStruct(
        foobar="foo",
        int=2,
        repeated=["value0", "value1"],
        nested=TestStruct(int=567),
        repeat_nested=[TestStruct(int=568)])
    expected_dict = {
        "foobar": "foo",
        "int": "2",  # Serialized
        "repeated": ["value0", "value1"],
        "nested": {
            "int": "567"  # Serialized
        },
        "repeat_nested": [{
            "int": "568"  # Serialized
        }]
    }
    self.assertEqual(
        test_struct.ToPrimitiveDict(stringify_leaf_fields=True), expected_dict)

  def testToPrimitiveDictStringifyBytes(self):
    data = b"\xff\xfe\xff"
    encoded_data = base64.b64encode(data).decode("ascii")

    class FooStruct(rdf_structs.RDFProtoStruct):

      type_description = type_info.TypeDescriptorSet(
          rdf_structs.ProtoBinary(name="data", field_number=1))

    foo_struct = FooStruct(data=data)
    foo_dict = foo_struct.ToPrimitiveDict(stringify_leaf_fields=True)

    self.assertIsInstance(foo_dict["data"], Text)
    self.assertEqual(foo_dict["data"], encoded_data)

  def testToPrimitiveDictStringifyRDFBytes(self):
    data = b"\xff\xfe\xff"
    encoded_data = base64.b64encode(data).decode("ascii")

    class BarStruct(rdf_structs.RDFProtoStruct):

      type_description = type_info.TypeDescriptorSet(
          rdf_structs.ProtoRDFValue(
              name="data", field_number=1, rdf_type=rdfvalue.RDFBytes))

    bar_struct = BarStruct(data=data)
    bar_dict = bar_struct.ToPrimitiveDict(stringify_leaf_fields=True)

    self.assertIsInstance(bar_dict["data"], Text)
    self.assertEqual(bar_dict["data"], encoded_data)

  def _GenerateSampleWithManyFields(self):
    fields = {}
    for _ in range(50):
      key = "foobar_%d" % random.randrange(100)
      fields[key] = key

    sample = TestStructWithManyFields(**fields)

    parsed = TestStructWithManyFields.FromSerializedBytes(
        sample.SerializeToBytes())

    return sample, parsed

  def testSerializationIsStable(self):
    sample1, sample2 = self._GenerateSampleWithManyFields()

    self.assertEqual(sample1.SerializeToBytes(), sample2.SerializeToBytes())

  def testRepeatedFieldHelperComparesToNone(self):
    self.assertNotEqual(TestStruct().repeated, None)

  def testSymmetricEqualityForDisjointFields(self):
    # Prevent a regression, where RDFStruct.__eq__ would be True for instances
    # having the same number of fields, but different fields which are set to
    # the default value.
    self.assertNotEqual(
        rdf_client_stats.IOSample(write_count=6),
        rdf_client_stats.IOSample(read_bytes=0))
    self.assertNotEqual(
        rdf_client_stats.IOSample(read_bytes=0),
        rdf_client_stats.IOSample(write_count=6))

  def testDefaultRepeatedSetterDoesNotChangeEquality(self):
    sample = TestStruct()
    sample.repeated  #  pylint: disable=pointless-statement
    self.assertEqual(sample, TestStruct())

  def testUnsetFieldsHaveSymmetricEqualityWithDefaultValues(self):
    self.assertEqual(TestStruct(repeated=[]), TestStruct())
    self.assertEqual(TestStruct(), TestStruct(repeated=[]))

  def testCanDeserializeOldVersion(self):
    a = VersionedTestStructV2(bool1=True, bool2=True)
    a_ser = a.SerializeToBytes()
    a_old = VersionedTestStructV1.FromSerializedBytes(a_ser)

    with self.assertRaises(AttributeError):
      a_old.Get("bool2")

  def testCanCompareEqualityAcrossVersions(self):
    a = VersionedTestStructV1.FromSerializedBytes(
        VersionedTestStructV2(bool1=True, bool2=True).SerializeToBytes())

    b = VersionedTestStructV1(bool1=True)

    self.assertEqual(a, b)
    self.assertEqual(b, a)


class BooleanToEnumMigrationTest(test_lib.GRRBaseTest):

  def testBooleansAreParsedAsEnums(self):
    a = TestStructWithBool()
    b = TestStructWithEnum.FromSerializedBytes(a.SerializeToBytes())
    self.assertEqual(b.foo, "ZERO")

    a = TestStructWithBool(foo=False)
    b = TestStructWithEnum.FromSerializedBytes(a.SerializeToBytes())
    self.assertEqual(b.foo, "ZERO")

    a = TestStructWithBool(foo=True)
    b = TestStructWithEnum.FromSerializedBytes(a.SerializeToBytes())
    self.assertEqual(b.foo, "ONE")

  def testEnumsAreParsedAsBooleans(self):
    a = TestStructWithEnum()
    b = TestStructWithBool.FromSerializedBytes(a.SerializeToBytes())
    self.assertIs(b.foo, False)

    a = TestStructWithEnum(foo="ZERO")
    b = TestStructWithBool.FromSerializedBytes(a.SerializeToBytes())
    self.assertIs(b.foo, False)

    a = TestStructWithEnum(foo="ONE")
    b = TestStructWithBool.FromSerializedBytes(a.SerializeToBytes())
    self.assertIs(b.foo, True)

  def testNewEnumOptionIsBackwardsCompatibleToTrue(self):
    a = TestStructWithEnum(foo="TWO")
    b = TestStructWithBool.FromSerializedBytes(a.SerializeToBytes())
    self.assertIs(b.foo, True)


class GenericRDFProtoTest(test_lib.GRRBaseTest):

  def testNestedProtobufAssignment(self):
    """Check that we can assign a nested protobuf."""
    container = rdf_client.BufferReference()
    test_path = "C:\\test"
    pathspec = rdf_paths.PathSpec(path=test_path, pathtype=1)

    # Should raise - incompatible RDFType.
    self.assertRaises(ValueError, setattr, container, "pathspec",
                      rdfvalue.RDFString("hello"))

    # Should raise - incompatible RDFProto type.
    self.assertRaises(ValueError, setattr, container, "pathspec",
                      rdf_client_fs.StatEntry(st_size=5))

    # Assign directly.
    container.device = pathspec

    self.assertEqual(container.device.path, test_path)

    # Clear the field.
    container.device = None

    # Check the protobuf does not have the field set at all.
    self.assertFalse(container.HasField("pathspec"))

  def testSimpleTypeAssignment(self):
    sample = rdf_client_fs.StatEntry()
    sample.AddDescriptor(
        rdf_structs.ProtoRDFValue(
            name="test",
            field_number=45,
            default=rdfvalue.RDFInteger(0),
            rdf_type=rdfvalue.RDFInteger))

    self.assertIsInstance(sample.test, rdfvalue.RDFInteger)

    # Can we assign an RDFValue instance?
    sample.test = rdfvalue.RDFInteger(5)

    self.assertEqual(sample.test, 5)

    # Check that bare values can be coerced.
    sample.test = 6
    self.assertIsInstance(sample.test, rdfvalue.RDFInteger)
    self.assertEqual(sample.test, 6)

    # Assign an enum.
    sample.registry_type = sample.RegistryType.REG_DWORD
    self.assertEqual(sample.registry_type, sample.RegistryType.REG_DWORD)

    sample.registry_type = rdf_client_fs.StatEntry.RegistryType.REG_SZ
    self.assertEqual(sample.registry_type, sample.RegistryType.REG_SZ)

    # We can also assign the string value.
    sample.registry_type = "REG_QWORD"
    self.assertEqual(sample.registry_type, sample.RegistryType.REG_QWORD)

    # Check that coercing works.
    sample.test = "10"
    self.assertEqual(sample.test, 10)

    # Assign an RDFValue which can not be coerced.
    self.assertRaises(type_info.TypeValueError, setattr, sample, "test",
                      rdfvalue.RDFString("hello"))

  def testComplexConstruction(self):
    """Test that we can construct RDFProtos with nested fields."""
    pathspec = rdf_paths.PathSpec(
        path="/foobar", pathtype=rdf_paths.PathSpec.PathType.TSK)
    sample = rdf_client_fs.StatEntry(pathspec=pathspec, st_size=5)

    self.assertEqual(sample.pathspec.path, "/foobar")
    self.assertEqual(sample.st_size, 5)

    self.assertRaises(AttributeError, rdf_client_fs.StatEntry, foobar=1)

  def testUnicodeSupport(self):
    pathspec = rdf_paths.PathSpec(
        path="/foobar", pathtype=rdf_paths.PathSpec.PathType.TSK)
    pathspec.path = u"Grüezi"

    self.assertEqual(pathspec.path, u"Grüezi")

  def testRepeatedFields(self):
    """Test handling of protobuf repeated fields."""
    sample = rdf_client_network.Interface()

    # Add an invalid type.
    self.assertRaises(type_info.TypeValueError, sample.addresses.Append, 2)

    # Add an rdfvalue by kwargs.
    sample.addresses.Append(human_readable_address="127.0.0.1")

    self.assertEqual(sample.addresses[0].human_readable_address, "127.0.0.1")
    self.assertLen(sample.addresses, 1)

    # Add an rdfvalue.
    sample.addresses.Append(
        rdf_client_network.NetworkAddress(human_readable_address="1.2.3.4"))

    self.assertLen(sample.addresses, 2)
    self.assertEqual(sample.addresses[1].human_readable_address, "1.2.3.4")

  def testEnums(self):
    """Check that enums are wrapped in a descriptor class."""
    sample = rdf_flows.GrrStatus()
    self.assertEqual(sample.status, rdf_flows.GrrStatus.ReturnedStatus.OK)


class EnumNamedValueTest(absltest.TestCase):

  def testInitialization(self):
    a = rdf_structs.EnumNamedValue(initializer=5, name="A", description="B")
    self.assertEqual(a.id, 5)
    self.assertEqual(a.name, "A")
    self.assertEqual(a.description, "B")

    b = rdf_structs.EnumNamedValue()
    self.assertEqual(b.id, 0)
    self.assertEqual(b.name, "0")
    self.assertIsNone(b.description)

  def testComparison(self):
    a = rdf_structs.EnumNamedValue(initializer=0, name="A")
    b = rdf_structs.EnumNamedValue(initializer=1, name="B")

    self.assertEqual(a, rdf_structs.EnumNamedValue(0, name="A"))
    self.assertNotEqual(a, b)
    self.assertLess(a, b)
    self.assertGreater(b, a)

    self.assertEqual(b, 1)
    self.assertEqual(b, "B")

  def testSerialization(self):
    # Current EnumNamedValue implementation only persists the ID.
    a = rdf_structs.EnumNamedValue(initializer=3)
    a2 = rdf_structs.EnumNamedValue.FromSerializedBytes(a.SerializeToBytes())

    self.assertEqual(a2.id, 3)
    self.assertEqual(a, a2)

    a2 = rdf_structs.EnumNamedValue.FromWireFormat(a.SerializeToWireFormat())
    self.assertEqual(a2.id, 3)
    self.assertEqual(a, a2)

  def testHumanReadable(self):
    a = rdf_structs.EnumNamedValue(initializer=5, name="A")
    self.assertEqual(str(a), "A")

    a2 = rdf_structs.EnumNamedValue.FromHumanReadable(str(a))
    self.assertEqual(a2.name, "A")
    self.assertEqual(a2, a)


class EnumContainerTest(absltest.TestCase):

  def setUp(self):
    super().setUp()
    self.enum_container = rdf_structs.EnumContainer(
        name="foo", values={
            "bar": 1,
            "baz": 2
        })

  def testFromString(self):
    self.assertEqual(
        self.enum_container.FromString("bar"), self.enum_container.bar)

  def testFromString_invalidValue(self):
    with self.assertRaises(ValueError):
      self.enum_container.FromString("bax")


class NoDynamicTypeLookupTest(absltest.TestCase):

  class NoDynamicTypeLookupMessage(rdf_structs.RDFProtoStruct):

    protobuf = tests_pb2.NoDynamicTypeLookupMessage

  def testFieldAccessors(self):
    message = self.NoDynamicTypeLookupMessage()
    self.assertIsInstance(message.any, rdf_structs.AnyValue)

    message.any.type_url = "foo.bar.baz"
    message.any.value = b"quux"

    self.assertEqual(message.any.type_url, "foo.bar.baz")
    self.assertEqual(message.any.value, b"quux")

  def testSerializeAndDeserialize(self):
    message = self.NoDynamicTypeLookupMessage()
    message.any.type_url = "foo.bar.baz"
    message.any.value = b"quux"

    blob = message.SerializeToBytes()
    message = self.NoDynamicTypeLookupMessage.FromSerializedBytes(blob)

    self.assertEqual(message.any.type_url, "foo.bar.baz")
    self.assertEqual(message.any.value, b"quux")

  def testAsPrimitiveProto(self):
    message = self.NoDynamicTypeLookupMessage()
    message.any.type_url = "foo.bar.baz"
    message.any.value = b"quux"

    proto = message.AsPrimitiveProto()
    self.assertEqual(proto.any.type_url, "foo.bar.baz")
    self.assertEqual(proto.any.value, b"quux")


class AnyValueTest(absltest.TestCase):

  def testFromProto2(self):
    string = wrappers_pb2.StringValue(value="Lorem ipsum.")

    string_proto2_any = any_pb2.Any()
    string_proto2_any.Pack(string)

    string_rdf_any = rdf_structs.AnyValue.FromProto2(string_proto2_any)

    unpacked_string = wrappers_pb2.StringValue()
    unpacked_string.ParseFromString(string_rdf_any.value)

    self.assertEqual(unpacked_string, string)

  def testPack(self):
    user = rdf_client.User()
    user.username = "foo"

    proto = rdf_structs.AnyValue.Pack(user)
    self.assertIn("User", proto.type_url)

    deserialized = rdf_client.User.FromSerializedBytes(proto.value)
    self.assertEqual(deserialized.username, "foo")

  def testUnpack(self):
    user = rdf_client.User()
    user.username = "foo"

    proto = rdf_structs.AnyValue.Pack(user)

    unpacked = proto.Unpack(rdf_client.User)
    self.assertEqual(unpacked.username, "foo")

  def testUnpackIncompatibleTypeURL(self):
    user = rdf_client.User()
    user.username = "foo"

    proto = rdf_structs.AnyValue.Pack(user)
    proto.type_url = "type.googleapis.com/com.example.Foo"

    unpacked = proto.Unpack(rdf_client.User)
    self.assertEqual(unpacked.username, "foo")

  def testPackProto2(self):
    string = wrappers_pb2.StringValue(value="Lorem ipsum.")

    packed_string = rdf_structs.AnyValue.PackProto2(string)
    self.assertIn("StringValue", packed_string.type_url)

    unpacked_string = wrappers_pb2.StringValue()
    unpacked_string.ParseFromString(packed_string.value)

    self.assertEqual(unpacked_string, string)

  def testPackProto2RDFValue(self):
    user = rdf_client.User()
    user.username = "foo"

    packed_user = rdf_structs.AnyValue.PackProto2(user.AsPrimitiveProto())

    unpacked_user = packed_user.Unpack(rdf_client.User)
    self.assertEqual(unpacked_user, user)


class VarintTest(absltest.TestCase):
  """Tests the VarintEncode and VarintReader implementations."""

  _ENCODED_VARINTS = (
      (1, b"\x01"),
      (1 << 4, b"\x10"),
      (1 << 8, b"\x80\x02"),
      (1 << 12, b"\x80\x20"),
      (1 << 16, b"\x80\x80\x04"),
      (1 << 62, b"\x80\x80\x80\x80\x80\x80\x80\x80\x40"),
      (1 << 63, b"\x80\x80\x80\x80\x80\x80\x80\x80\x80\x01"),
  )

  def testCorrectlyEncodesIntegers(self):
    for val, b in self._ENCODED_VARINTS:
      with self.subTest(value=val):
        self.assertEqual(rdf_structs.VarintEncode(val), b)

  def testEncodesOverflowIntegerAsZero(self):
    self.assertEqual(rdf_structs.VarintEncode(1 << 66), b"\x00")

  def testCorrectlyReadsEncodedIntegers(self):
    for src_val, b in self._ENCODED_VARINTS:
      with self.subTest(buffer=b):
        self.assertEqual(rdf_structs.VarintReader(b, 0), (src_val, len(b)))

  def testReaderCanReadWhatEncoderHasEncoded(self):
    for v in [
        0, 1, 2, 10, 20, 63, 64, 65, 127, 128, 129, 255, 256, 257, 1 << 63 - 1,
        1 << 64 - 1
    ]:
      with self.subTest(value=v):
        buf = rdf_structs.VarintEncode(v)
        r, p = rdf_structs.VarintReader(buf, 0)
        self.assertEqual(v, r)
        self.assertLen(buf, p)

  def testDecodingZeroBufferRaises(self):
    with self.assertRaisesRegex(ValueError,
                                "Too many bytes when decoding varint"):
      rdf_structs.VarintReader(b"", 0)

  def testDecodingValueLargerThan64BitReturnsTruncatedValue(self):
    self.assertEqual(
        rdf_structs.VarintReader(b"\x80\x80\x80\x80\x80\x80\x80\x80\x80\x02",
                                 0), (0, 10))

  def testRaisesWhenBufferIsTooLong(self):
    with self.assertRaisesRegex(ValueError,
                                "Too many bytes when decoding varint"):
      rdf_structs.VarintReader(
          b"\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x80\x01\x00\x00", 0)
    with self.assertRaisesRegex(ValueError,
                                "Too many bytes when decoding varint"):
      rdf_structs.VarintReader(b"\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF",
                               0)

  def testNonCanonicalZeroDoesNotRaise(self):
    self.assertEqual(rdf_structs.VarintReader(b"\x80\x80\x80\x00", 0), (0, 4))


class SignedVarintTest(absltest.TestCase):
  """Test the VarintReader implementation."""

  _ENCODED_SIGNED_VARINTS = (
      (-1, b"\xff\xff\xff\xff\xff\xff\xff\xff\xff\x01"),
      (-(1 << 4), b"\xf0\xff\xff\xff\xff\xff\xff\xff\xff\x01"),
      (-(1 << 8), b"\x80\xfe\xff\xff\xff\xff\xff\xff\xff\x01"),
      (-(1 << 12), b"\x80\xe0\xff\xff\xff\xff\xff\xff\xff\x01"),
      (-(1 << 16), b"\x80\x80\xfc\xff\xff\xff\xff\xff\xff\x01"),
      (-(1 << 62), b"\x80\x80\x80\x80\x80\x80\x80\x80\xc0\x01"),
      (-(1 << 63), b"\x80\x80\x80\x80\x80\x80\x80\x80\x80\x01"),
  )

  def testCorrectlyEncodesSignedIntegers(self):
    for val, b in self._ENCODED_SIGNED_VARINTS:
      with self.subTest(value=val):
        self.assertEqual(rdf_structs.SignedVarintEncode(val), b)

  def testEncodesOverflowIntegerAsZero(self):
    self.assertEqual(rdf_structs.SignedVarintEncode(-(1 << 66)), b"\x00")

  def testCorrectlyReadsEncodedSignedIntegers(self):
    for src_val, b in self._ENCODED_SIGNED_VARINTS:
      with self.subTest(buffer=b):
        self.assertEqual(
            rdf_structs.SignedVarintReader(b, 0), (src_val, len(b)))


class SplitBufferTest(absltest.TestCase):
  """Test the SplitBuffer function."""

  def testRaisesIfIndexOrLengthIncorrect(self):
    buf = b"\x00\x00"
    with self.assertRaises(ValueError):
      rdf_structs.SplitBuffer(buf, -1, 2)
    with self.assertRaises(ValueError):
      rdf_structs.SplitBuffer(buf, 0, -1)
    with self.assertRaises(ValueError):
      rdf_structs.SplitBuffer(buf, -1, -1)
    with self.assertRaises(ValueError):
      rdf_structs.SplitBuffer(buf, 3, 2)

  def testReturnsNothingOnEmptyBuffer(self):
    self.assertEqual(rdf_structs.SplitBuffer(b"", 0, 0), [])
    self.assertEqual(rdf_structs.SplitBuffer(b"\x00\x00", 2, 2), [])

  def testRaisesOnBrokenTag(self):
    with self.assertRaisesRegex(ValueError, "Broken tag encountered"):
      rdf_structs.SplitBuffer(b"\xff", 0, 1)

  def testRaisesOnBrokenVarintTag(self):
    # This used to trigger an infinite loop in the accelerated.c implementation.
    with self.assertRaisesRegex(ValueError, "Broken varint tag encountered"):
      rdf_structs.SplitBuffer(b"\x00", 0, 1)

    with self.assertRaisesRegex(ValueError, "Broken varint tag encountered"):
      rdf_structs.SplitBuffer(b"\x00\xff", 0, 2)

  def testCorrectlyProcessesVarintTag(self):
    self.assertEqual(
        rdf_structs.SplitBuffer("\x00\x01", 0, 2), [(b"\x00", b"", b"\x01")])

  def testRaisesOnOversizedFixed64Tag(self):
    for l in range(8):
      with self.subTest(l=l):
        with self.assertRaisesRegex(ValueError,
                                    "Fixed64 tag exceeds available buffer"):
          buf = b"\x01" + b"\x00" * l
          rdf_structs.SplitBuffer(buf, 0, len(buf))

  def testCorrectlyProcessesFixed64Tag(self):
    self.assertEqual(
        rdf_structs.SplitBuffer("\x01\x00\x01\x02\x03\x04\x05\0x6\0x7", 0, 9),
        [(b"\x01", b"", b"\x00\x01\x02\x03\x04\x05\x00x")])

  def testRaisesOnOversizedFixed32Tag(self):
    for l in range(4):
      with self.subTest(l=l):
        with self.assertRaisesRegex(ValueError,
                                    "Fixed32 tag exceeds available buffer"):
          buf = b"\x05" + b"\x00" * l
          rdf_structs.SplitBuffer(buf, 0, len(buf))

  def testCorrectlyProcessesFixed32Tag(self):
    self.assertEqual(
        rdf_structs.SplitBuffer("\x05\x00\x01\x02\x03", 0, 5),
        [(b"\x05", b"", b"\x00\x01\x02\x03")])

  def testRaisesOnBrokenLengthDelimitedTag(self):
    with self.assertRaisesRegex(ValueError,
                                "Broken length_delimited tag encountered"):
      rdf_structs.SplitBuffer(b"\x02\xff", 0, 2)

  def testRaisesOnLengthDelimitedTagExceedingMaxInt(self):
    with self.assertRaisesRegex(ValueError, "Length delimited exceeds limits"):
      buf = b"\x0a\xff\xff\xff\xff\xff\xff\xff\xff\xff\x01\x02\xff\xff\xff\xff\xff\xff\x27"
      rdf_structs.SplitBuffer(buf, 0, len(buf))

  def testRaisesOnOversizedLengthDelimitedTag(self):
    with self.assertRaisesRegex(ValueError,
                                "Length tag exceeds available buffer"):
      rdf_structs.SplitBuffer(b"\x02\x02", 0, 2)

  def testCorrectlyProcessesLengthDelimitedTag(self):
    self.assertEqual(
        rdf_structs.SplitBuffer("\x02\x02\x00\x01", 0, 4),
        [(b"\x02", b"\x02", b"\x00\x01")])

  def testRaisesOnUnknownTag(self):
    with self.assertRaisesRegex(ValueError, "Unexpected Tag"):
      rdf_structs.SplitBuffer(b"\x03", 0, 1)


def main(argv):
  test_lib.main(argv)


if __name__ == "__main__":
  app.run(main)
