1 files changed, 4 insertions, 7 deletions

diff --git a/tests/sys/common/divert.py b/tests/sys/common/divert.py
index 94e20a03571f..061f787d1b7a 100755
--- a/tests/sys/common/divert.py
+++ b/tests/sys/common/divert.py
@@ -29,16 +29,13 @@
 #
 
 
-import socket
+from socket import socket, PF_DIVERT, SOCK_RAW
 import logging
 logging.getLogger("scapy").setLevel(logging.CRITICAL)
 import scapy.all as sc
 import argparse
 
 
-IPPROTO_DIVERT = 258
-
-
 def parse_args():
     parser = argparse.ArgumentParser(description='divert socket tester')
     parser.add_argument('--dip', type=str, help='destination packet IP')
@@ -52,14 +49,14 @@ def parse_args():
 
 def ipdivert_ip_output_remote_success(args):
     packet = sc.IP(dst=args.dip) / sc.ICMP(type='echo-request')
-    with socket.socket(socket.AF_INET, socket.SOCK_RAW, IPPROTO_DIVERT) as s:
+    with socket(PF_DIVERT, SOCK_RAW, 0) as s:
         s.bind(('0.0.0.0', args.divert_port))
         s.sendto(bytes(packet), ('0.0.0.0', 0))
 
 
 def ipdivert_ip6_output_remote_success(args):
     packet = sc.IPv6(dst=args.dip) / sc.ICMPv6EchoRequest()
-    with socket.socket(socket.AF_INET, socket.SOCK_RAW, IPPROTO_DIVERT) as s:
+    with socket(PF_DIVERT, SOCK_RAW, 0) as s:
         s.bind(('0.0.0.0', args.divert_port))
         s.sendto(bytes(packet), ('0.0.0.0', 0))
 
@@ -67,7 +64,7 @@ def ipdivert_ip6_output_remote_success(args):
 def ipdivert_ip_input_local_success(args):
     """Sends IPv4 packet to OS stack as inbound local packet."""
     packet = sc.IP(dst=args.dip,src=args.sip) / sc.ICMP(type='echo-request')
-    with socket.socket(socket.AF_INET, socket.SOCK_RAW, IPPROTO_DIVERT) as s:
+    with socket(PF_DIVERT, SOCK_RAW, 0) as s:
         s.bind(('0.0.0.0', args.divert_port))
         s.sendto(bytes(packet), (args.dip, 0))