Add some more tests and fix comments in IPv4 class

Changes due to PR review on github.

Author: renanrodrigo

pyof/foundation/network_types.py

@@ -145,33 +145,38 @@ class IPv4(GenericStruct):
     Contains all fields of an IP version 4 packet header, plus the upper layer
     content as binary data.
     Some of the fields were merged together because of their size being
-    inferior to 8 bits.
+    inferior to 8 bits. They are represented as a single class attribute, but
+    pack/unpack methods will take into account the values in individual
+    instance attributes.
     """
 
-    # IP protocol version + Internet Header Length (words)
+    #: _version_ihl (:class:`UBInt8`): IP protocol version + Internet Header
+    #: Length (words)
     _version_ihl = UBInt8()
-    # Differentiated Services Code Point (ToS - Type of Service) +
-    # Explicit Congestion Notification
+    #: _dscp_ecn (:class:`UBInt8`): Differentiated Services Code Point
+    #: (ToS - Type of Service) + Explicit Congestion Notification
     _dscp_ecn = UBInt8()
-    # IP packet length (bytes)
+    #: length (:class:`UBInt16`): IP packet length (bytes)
     length = UBInt16()
-    # Packet ID - common to all fragments
+    #: identification (:class:`UBInt16`): Packet ID - common to all fragments
     identification = UBInt16()
-    # Fragmentation flags + fragmentation offset
+    #: _flags_offset (:class:`UBInt16`): Fragmentation flags + fragmentation
+    #: offset
     _flags_offset = UBInt16()
-    # Packet time-to-live
+    #: ttl (:class:`UBInt8`): Packet time-to-live
     ttl = UBInt8()
-    # Upper layer protocol number
+    #: protocol (:class:`UBInt8`): Upper layer protocol number
     protocol = UBInt8()
-    # Header checksum
+    #: checksum (:class:`UBInt16`): Header checksum
     checksum = UBInt16()
-    # Source address
+    #: source (:class:`IPAddress`): Source IPv4 address
     source = IPAddress()
-    # Destination address
+    #: destination (:class:`IPAddress`): Destination IPv4 address
     destination = IPAddress()
-    # IP Options - up to 320 bits, always padded to 32 bits
+    #: options (:class:`BinaryData`): IP Options - up to 320 bits, always
+    #: padded to 32 bits
     options = BinaryData()
-    # Packet data
+    #: data (:class:`BinaryData`): Packet data
     data = BinaryData()
 
     def __init__(self, version=4, ihl=5, dscp=0, ecn=0, length=0, # noqa

tests/test_foundation/test_network_types.py

@@ -16,11 +16,44 @@ def test_GenTLV_value_unpack(self):
         tlv_unpacked.unpack(tlv.pack())
         self.assertEqual(tlv.value.value, tlv_unpacked.value.value)
 
-    def test_IPv4_pack_unpack(self):
+
+class TestIPv4(unittest.TestCase):
+    """Test IPv4 packets."""
+
+    def test_IPv4_pack(self):
         """Test pack/unpack of IPv4 class."""
-        packet = IPv4(ttl=64, protocol=17, source="192.168.0.1",
-                      destination="172.16.200.132", data=b'testdata')
+        packet = IPv4(dscp=10, ttl=64, protocol=17, source="192.168.0.10",
+                      destination="172.16.10.30", options=b'1000',
+                      data=b'testdata')
         packed = packet.pack()
+        expected = b'F(\x00 \x00\x00\x00\x00@\x11\x02'
+        expected += b'\xc5\xc0\xa8\x00\n\xac\x10\n\x1e1000testdata'
+        self.assertEqual(packed, expected)
+
+    def test_IPv4_unpack(self):
+        """Test unpack of IPv4 binary packet."""
+        raw = b'FP\x00$\x00\x00\x00\x00\x80\x06W'
+        raw += b'\xf4\n\x9aN\x81\xc0\xa8\xc7\xcc1000somemoredata'
+        expected = IPv4(dscp=20, ttl=128, protocol=6, source="10.154.78.129",
+                        destination="192.168.199.204", options=b'1000',
+                        data=b'somemoredata')
+        expected.pack()
         unpacked = IPv4()
-        unpacked.unpack(packed)
-        self.assertEqual(packed, unpacked.pack())
+        unpacked.unpack(raw)
+        self.assertEqual(unpacked, expected)
+
+    def test_IPv4_size(self):
+        """Test Header size for IPv4 packet."""
+        packet = IPv4()
+        packet.pack()
+        self.assertEqual(20, packet.get_size())
+        self.assertEqual(20, packet.length)
+        self.assertEqual(20, packet.ihl * 4)
+
+    def test_IPv4_checksum(self):
+        """Test if the IPv4 checksum is being calculated correclty."""
+        packet = IPv4(dscp=10, ttl=64, protocol=17, source="192.168.0.10",
+                      destination="172.16.10.30", options=b'1000',
+                      data=b'testdata')
+        packet.pack()
+        self.assertEqual(packet.checksum, 709)