[PATCH v2 18/22] tasst: IP address allocation helpers

[David Gibson <david@gibson.dropbear.id.au>]

A bunch of our test scenarious will require us to allocate IPv4 and IPv6
addresses in example networks.  Make helpers to do this easily.

Signed-off-by: David Gibson <david@gibson.dropbear.id.au>
---
 test/Makefile               |  2 +-
 test/tasst/address.py       | 79 +++++++++++++++++++++++++++++++++++++
 test/tasst/selftest/veth.py | 41 ++++++++++++++++++-
 test/tasst/transfer.py      |  6 +--
 4 files changed, 123 insertions(+), 5 deletions(-)
 create mode 100644 test/tasst/address.py

diff --git a/test/Makefile b/test/Makefile
index 584f56e9..f3a3cc58 100644
--- a/test/Makefile
+++ b/test/Makefile
@@ -72,7 +72,7 @@ EXETER_PY = build/build.py
 EXETER_JOBS = $(EXETER_SH:%.sh=%.json) $(EXETER_PY:%.py=%.json)
 AVOCADO_JOBS = $(EXETER_JOBS) avocado/static_checkers.json
 
-TASST_SRCS = __init__.py __main__.py nstool.py snh.py transfer.py \
+TASST_SRCS = __init__.py __main__.py address.py nstool.py snh.py transfer.py \
 	selftest/__init__.py selftest/static_ifup.py selftest/veth.py
 
 EXETER_META = meta/lint.json meta/tasst.json
diff --git a/test/tasst/address.py b/test/tasst/address.py
new file mode 100644
index 00000000..70899789
--- /dev/null
+++ b/test/tasst/address.py
@@ -0,0 +1,79 @@
+#! /usr/bin/env python3
+
+# SPDX-License-Identifier: GPL-2.0-or-later
+#
+# Copyright Red Hat
+# Author: David Gibson <david@gibson.dropbear.id.au>
+
+"""
+Test A Simple Socket Transport
+
+address.py - Address allocation helpers
+"""
+
+import ipaddress
+
+import exeter
+
+# Loopback addresses, for convenience
+LOOPBACK4 = ipaddress.ip_address('127.0.0.1')
+LOOPBACK6 = ipaddress.ip_address('::1')
+
+# Documentation test networks defined by RFC 5737
+TEST_NET_1 = ipaddress.ip_network('192.0.2.0/24')
+TEST_NET_2 = ipaddress.ip_network('198.51.100.0/24')
+TEST_NET_3 = ipaddress.ip_network('203.0.113.0/24')
+
+# Documentation test network defined by RFC 3849
+TEST_NET6 = ipaddress.ip_network('2001:db8::/32')
+# Some subnets of that for our usage
+TEST_NET6_TASST_A = ipaddress.ip_network('2001:db8:9a55:aaaa::/64')
+TEST_NET6_TASST_B = ipaddress.ip_network('2001:db8:9a55:bbbb::/64')
+TEST_NET6_TASST_C = ipaddress.ip_network('2001:db8:9a55:cccc::/64')
+
+
+class IpiAllocator:
+    """IP address allocator"""
+
+    DEFAULT_NETS = [TEST_NET_1, TEST_NET6_TASST_A]
+
+    def __init__(self, *nets):
+        if not nets:
+            nets = self.DEFAULT_NETS
+
+        self.nets = [ipaddress.ip_network(n) for n in nets]
+        self.hostses = [n.hosts() for n in self.nets]
+
+    def next_ipis(self):
+        addrs = [next(h) for h in self.hostses]
+        return [ipaddress.ip_interface(f'{a}/{n.prefixlen}')
+                for a, n in zip(addrs, self.nets)]
+
+
+@exeter.test
+def ipa_test(nets=None, count=12):
+    if nets is None:
+        ipa = IpiAllocator()
+        nets = IpiAllocator.DEFAULT_NETS
+    else:
+        ipa = IpiAllocator(*nets)
+
+    addrsets = [set() for i in range(len(nets))]
+    for i in range(count):
+        addrs = ipa.next_ipis()
+        # Check we got as many addresses as expected
+        exeter.assert_eq(len(addrs), len(nets))
+        for s, a, n in zip(addrsets, addrs, nets):
+            # Check the addresses belong to the right network
+            exeter.assert_eq(a.network, ipaddress.ip_network(n))
+            s.add(a)
+
+    print(addrsets)
+    # Check the addresses are unique
+    for s in addrsets:
+        exeter.assert_eq(len(s), count)
+
+
+@exeter.test
+def ipa_test_custom():
+    ipa_test(nets=['10.55.0.0/16', '192.168.55.0/24', 'fd00:9a57:a000::/48'])
diff --git a/test/tasst/selftest/veth.py b/test/tasst/selftest/veth.py
index 24bbdc27..39ac947d 100644
--- a/test/tasst/selftest/veth.py
+++ b/test/tasst/selftest/veth.py
@@ -16,7 +16,7 @@ import ipaddress
 
 import exeter
 
-from tasst import nstool
+from tasst import address, nstool, transfer
 
 
 @contextlib.contextmanager
@@ -65,3 +65,42 @@ def test_optimistic_dad():
 @exeter.test
 def test_no_dad():
     test_slaac(dad='disable')
+
+
+@contextlib.contextmanager
+def configured_veth(ip1, ip2):
+    with unconfigured_veth() as (ns1, ns2):
+        ns1.ifup('lo')
+        ns1.ifup('veth1', ip1, dad='disable')
+
+        ns2.ifup('lo')
+        ns2.ifup('veth2', ip2, dad='disable')
+
+        yield (ns1, ns2)
+
+
+@contextlib.contextmanager
+def veth_transfer(ip1, ip2):
+    with configured_veth(ip1, ip2) as (ns1, ns2):
+        yield transfer.TransferTestScenario(client=ns1, server=ns2,
+                                            connect_ip=ip2.ip,
+                                            connect_port=10000)
+
+
+ipa = address.IpiAllocator()
+NS1_IP4, NS1_IP6 = ipa.next_ipis()
+NS2_IP4, NS2_IP6 = ipa.next_ipis()
+
+
+def veth_transfer4():
+    return veth_transfer(NS1_IP4, NS2_IP4)
+
+
+transfer.transfer_tests(veth_transfer4)
+
+
+def veth_transfer6():
+    return veth_transfer(NS1_IP6, NS2_IP6)
+
+
+transfer.transfer_tests(veth_transfer6)
diff --git a/test/tasst/transfer.py b/test/tasst/transfer.py
index be3eebc2..a5aa0614 100644
--- a/test/tasst/transfer.py
+++ b/test/tasst/transfer.py
@@ -17,7 +17,7 @@ import time
 
 import exeter
 
-from . import nstool, snh
+from . import address, nstool, snh
 
 
 # HACK: how long to wait for the server to be ready and listening (s)
@@ -175,7 +175,7 @@ def local_transfer4():
     with nstool.unshare_snh('ns', '-Un') as ns:
         ns.ifup('lo')
         yield TransferTestScenario(client=ns, server=ns,
-                                   connect_ip=IPv4Address('127.0.0.1'),
+                                   connect_ip=address.LOOPBACK4,
                                    connect_port=10000)
 
 
@@ -187,7 +187,7 @@ def local_transfer6():
     with nstool.unshare_snh('ns', '-Un') as ns:
         ns.ifup('lo')
         yield TransferTestScenario(client=ns, server=ns,
-                                   connect_ip=IPv6Address('::1'),
+                                   connect_ip=address.LOOPBACK6,
                                    connect_port=10000)
 
 
-- 
@@ -17,7 +17,7 @@ import time
 
 import exeter
 
-from . import nstool, snh
+from . import address, nstool, snh
 
 
 # HACK: how long to wait for the server to be ready and listening (s)
@@ -175,7 +175,7 @@ def local_transfer4():
     with nstool.unshare_snh('ns', '-Un') as ns:
         ns.ifup('lo')
         yield TransferTestScenario(client=ns, server=ns,
-                                   connect_ip=IPv4Address('127.0.0.1'),
+                                   connect_ip=address.LOOPBACK4,
                                    connect_port=10000)
 
 
@@ -187,7 +187,7 @@ def local_transfer6():
     with nstool.unshare_snh('ns', '-Un') as ns:
         ns.ifup('lo')
         yield TransferTestScenario(client=ns, server=ns,
-                                   connect_ip=IPv6Address('::1'),
+                                   connect_ip=address.LOOPBACK6,
                                    connect_port=10000)
 
 
-- 
2.45.2