Home | History | Annotate | Download | only in test 
      1 #!/usr/bin/python
      2 #
      3 # Copyright 2016 The Android Open Source Project
      4 #
      5 # Licensed under the Apache License, Version 2.0 (the "License");
      6 # you may not use this file except in compliance with the License.
      7 # You may obtain a copy of the License at
      8 #
      9 # http://www.apache.org/licenses/LICENSE-2.0
     10 #
     11 # Unless required by applicable law or agreed to in writing, software
     12 # distributed under the License is distributed on an "AS IS" BASIS,
     13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     14 # See the License for the specific language governing permissions and
     15 # limitations under the License.
     16 
     17 """Unit tests for csocket."""
     18 
     19 from socket import *  # pylint: disable=wildcard-import
     20 import unittest
     21 
     22 import csocket
     23 
     24 
     25 LOOPBACK_IFINDEX = 1
     26 SOL_IPV6 = 41
     27 
     28 
     29 class CsocketTest(unittest.TestCase):
     30 
     31   def _BuildSocket(self, family, addr):
     32     s = socket(family, SOCK_DGRAM, 0)
     33     s.bind((addr, 0))
     34     return s
     35 
     36   def CheckRecvfrom(self, family, addr):
     37     s = self._BuildSocket(family, addr)
     38     addr = s.getsockname()
     39     sockaddr = csocket.Sockaddr(addr)
     40     s.sendto("foo", addr)
     41     data, addr = csocket.Recvfrom(s, 4096, 0)
     42     self.assertEqual("foo", data)
     43     self.assertEqual(sockaddr, addr)
     44 
     45     s.close()
     46 
     47   def testRecvfrom(self):
     48     self.CheckRecvfrom(AF_INET, "127.0.0.1")
     49     self.CheckRecvfrom(AF_INET6, "::1")
     50 
     51   def CheckRecvmsg(self, family, addr):
     52     s = self._BuildSocket(family, addr)
     53 
     54     if family == AF_INET:
     55       s.setsockopt(SOL_IP, csocket.IP_PKTINFO, 1)
     56       s.setsockopt(SOL_IP, csocket.IP_RECVTTL, 1)
     57       pktinfo_addr = inet_pton(AF_INET, addr)
     58       pktinfo = (SOL_IP, csocket.IP_PKTINFO,
     59                  csocket.InPktinfo((LOOPBACK_IFINDEX,
     60                                     pktinfo_addr, pktinfo_addr)))
     61       ttl = (SOL_IP, csocket.IP_TTL, 64)
     62     elif family == AF_INET6:
     63       s.setsockopt(SOL_IPV6, csocket.IPV6_RECVPKTINFO, 1)
     64       s.setsockopt(SOL_IPV6, csocket.IPV6_RECVHOPLIMIT, 1)
     65       pktinfo_addr = inet_pton(AF_INET6, addr)
     66       pktinfo = (SOL_IPV6, csocket.IPV6_PKTINFO,
     67                  csocket.In6Pktinfo((pktinfo_addr, LOOPBACK_IFINDEX)))
     68       ttl = (SOL_IPV6, csocket.IPV6_HOPLIMIT, 64)
     69 
     70     addr = s.getsockname()
     71     sockaddr = csocket.Sockaddr(addr)
     72     s.sendto("foo", addr)
     73     data, addr, cmsg = csocket.Recvmsg(s, 4096, 1024, 0)
     74     self.assertEqual("foo", data)
     75     self.assertEqual(sockaddr, addr)
     76     self.assertEqual([pktinfo, ttl], cmsg)
     77 
     78     s.close()
     79 
     80   def testRecvmsg(self):
     81     self.CheckRecvmsg(AF_INET, "127.0.0.1")
     82     self.CheckRecvmsg(AF_INET6, "::1")
     83 
     84 
     85 if __name__ == "__main__":
     86   unittest.main()
     87