From a798d5ea777e20bb6b75e6a616e5eff7268578a7 Mon Sep 17 00:00:00 2001
From: Lorenzo Colitti <lorenzo@google.com>
Date: Wed, 1 Nov 2017 14:51:11 +0900
Subject: Provide an AddressVersion utility function.

Test: multinetwork_test.py passes on android-4.4
Test: neighbour_test.py passes on android-4.4
Change-Id: Id127ff72503a23ba4b42ee57699db5e967b4fad5
---
 net/test/csocket.py        | 4 ++++
 net/test/iproute.py        | 6 +++---
 net/test/neighbour_test.py | 7 ++++---
 3 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/net/test/csocket.py b/net/test/csocket.py
index 1dd58f1..4eff280 100644
--- a/net/test/csocket.py
+++ b/net/test/csocket.py
@@ -86,6 +86,10 @@ def LinuxVersion():
   return version
 
 
+def AddressVersion(addr):
+  return 6 if ":" in addr else 4
+
+
 def SetSocketTimeout(sock, ms):
   s = ms / 1000
   us = (ms % 1000) * 1000
diff --git a/net/test/iproute.py b/net/test/iproute.py
index cf5c55d..ba00f49 100644
--- a/net/test/iproute.py
+++ b/net/test/iproute.py
@@ -487,7 +487,7 @@ class IPRoute(netlink.NetlinkSocket):
       AssertionError: An IPv6 address was requested, and it did not appear
         within the timeout.
     """
-    version = 6 if ":" in address else 4
+    version = csocket.AddressVersion(address)
 
     flags = IFA_F_PERMANENT
     if version == 6:
@@ -502,7 +502,7 @@ class IPRoute(netlink.NetlinkSocket):
       self._WaitForAddress(sock, address, ifindex)
 
   def DelAddress(self, address, prefixlen, ifindex):
-    self._Address(6 if ":" in address else 4,
+    self._Address(csocket.AddressVersion(address),
                   RTM_DELADDR, address, prefixlen, 0, 0, ifindex)
 
   def GetAddress(self, address, ifindex=0):
@@ -553,7 +553,7 @@ class IPRoute(netlink.NetlinkSocket):
                 nexthop, dev, None, None)
 
   def GetRoutes(self, dest, oif, mark, uid, iif=None):
-    version = 6 if ":" in dest else 4
+    version = csocket.AddressVersion(dest)
     prefixlen = {4: 32, 6: 128}[version]
     self._Route(version, RTPROT_STATIC, RTM_GETROUTE, 0, dest, prefixlen, None,
                 oif, mark, uid, iif=iif)
diff --git a/net/test/neighbour_test.py b/net/test/neighbour_test.py
index 35b3b61..2caba44 100755
--- a/net/test/neighbour_test.py
+++ b/net/test/neighbour_test.py
@@ -22,6 +22,7 @@ import unittest
 
 from scapy import all as scapy
 
+import csocket
 import multinetwork_base
 import net_test
 
@@ -87,7 +88,7 @@ class NeighbourTest(multinetwork_base.MultiNetworkBaseTest):
     self.ifindex = self.ifindices[self.netid]
 
   def GetNeighbour(self, addr):
-    version = 6 if ":" in addr else 4
+    version = csocket.AddressVersion(addr)
     for msg, args in self.iproute.DumpNeighbours(version):
       if args["NDA_DST"] == addr:
         return msg, args
@@ -114,7 +115,7 @@ class NeighbourTest(multinetwork_base.MultiNetworkBaseTest):
         self.assertEquals(attrs[name], actual_attrs[name])
 
   def ExpectProbe(self, is_unicast, addr):
-    version = 6 if ":" in addr else 4
+    version = csocket.AddressVersion(addr)
     if version == 6:
       llsrc = self.MyMacAddress(self.netid)
       if is_unicast:
@@ -143,7 +144,7 @@ class NeighbourTest(multinetwork_base.MultiNetworkBaseTest):
 
   def ReceiveUnicastAdvertisement(self, addr, mac, srcaddr=None, dstaddr=None,
                                   S=1, O=0, R=1):
-    version = 6 if ":" in addr else 4
+    version = csocket.AddressVersion(addr)
     if srcaddr is None:
       srcaddr = addr
     if dstaddr is None:
-- 
cgit v1.2.3