[PATCH net-next 9/9] selftest/net: psp: Add test for dev-assoc/disassoc

[Wei Wang <weibunny@fb.com>]

Add an env NetDrvContEnPsp which is based on NetDrvContEn, but add an
additional bpf redirect program on nk_host to redirect traffic to the
psp_dev_local. The topology looks like this:
  Host NS:  psp_dev_local <---> nk_host
                |                |
                |                | (netkit pair)
                |                |
  Remote NS: psp_dev_peer      Guest NS: nk_guest
             (responder)             (PSP tests)

Add following tests for dev-assoc/dev-disassoc functionality:
1. Test the output of `./tools/net/ynl/pyynl/cli.py --spec
Documentation/netlink/specs/psp.yaml --dump dev-get` in both default and
the guest netns.
2. Test the case where we associate netkit with psp_dev_local, and
send PSP traffic from nk_guest to psp_dev_peer in 2 different netns.
3. Test to make sure the key rotation notification is sent to the netns
for associated dev as well
4. Test to make sure the dev change notification is sent to the netns
for associated dev as well
5. Test the deletion of nk_guest in client netns, and proper cleanup in
the assoc-list for psp dev.

Signed-off-by: Wei Wang <weibunny@fb.com>
---
 .../selftests/drivers/net/lib/py/__init__.py  |   6 +-
 .../selftests/drivers/net/lib/py/env.py       | 184 +++++++
 tools/testing/selftests/drivers/net/psp.py    | 462 +++++++++++++++++-
 3 files changed, 644 insertions(+), 8 deletions(-)

diff --git a/tools/testing/selftests/drivers/net/lib/py/__init__.py b/tools/testing/selftests/drivers/net/lib/py/__init__.py
index 6b55068d5370..ec512a9b8a2a 100644
--- a/tools/testing/selftests/drivers/net/lib/py/__init__.py
+++ b/tools/testing/selftests/drivers/net/lib/py/__init__.py
@@ -44,12 +44,12 @@ try:
                "ksft_ne", "ksft_not_in", "ksft_raises", "ksft_true", "ksft_gt",
                "ksft_not_none", "ksft_not_none"]
 
-    from .env import NetDrvEnv, NetDrvEpEnv, NetDrvContEnv
+    from .env import NetDrvEnv, NetDrvEpEnv, NetDrvContEnv, NetDrvContEnvPsp
     from .load import GenerateTraffic, Iperf3Runner
     from .remote import Remote
 
-    __all__ += ["NetDrvEnv", "NetDrvEpEnv", "NetDrvContEnv", "GenerateTraffic",
-                "Remote", "Iperf3Runner"]
+    __all__ += ["NetDrvEnv", "NetDrvEpEnv", "NetDrvContEnv", "NetDrvContEnvPsp",
+                "GenerateTraffic", "Remote", "Iperf3Runner"]
 except ModuleNotFoundError as e:
     print("Failed importing `net` library from kernel sources")
     print(str(e))
diff --git a/tools/testing/selftests/drivers/net/lib/py/env.py b/tools/testing/selftests/drivers/net/lib/py/env.py
index 857ae0f37516..83a4d79c6053 100644
--- a/tools/testing/selftests/drivers/net/lib/py/env.py
+++ b/tools/testing/selftests/drivers/net/lib/py/env.py
@@ -4,6 +4,7 @@ import ipaddress
 import os
 import time
 import json
+import re
 from pathlib import Path
 from lib.py import KsftSkipEx, KsftXfailEx
 from lib.py import ksft_setup, wait_file
@@ -452,3 +453,186 @@ class NetDrvContEnv(NetDrvEpEnv):
         value = ipv6_bytes + ifindex_bytes
         value_hex = ' '.join(f'{b:02x}' for b in value)
         bpftool(f"map update id {bss_map_id} key hex 00 00 00 00 value hex {value_hex}")
+
+class NetDrvContEnvPsp(NetDrvContEnv):
+    """
+    Class for PSP testing with netkit pair + BPF redirect topology.
+
+    Inherits from NetDrvContEnv and adds PSP-specific attributes.
+
+    Topology (3 namespaces):
+        Host Namespace:
+            psp_dev_local (2001:db8::1) ←──linked──→ psp_dev_peer (in _netns)
+            (PSP device)
+                 │
+                 │ BPF on nsim_local ingress: bpf_redirect_peer() to nk_guest
+                 │
+            nk_host (fe80::1)
+                 │
+                 │ BPF on nk_host ingress: bpf_redirect_neigh() to nsim_local
+                 │
+                 │ netkit pair
+                 │
+        Guest Namespace (netns):
+                 │
+            nk_guest (fe80::2, 2001:db9::2:2)
+            ★ PSP tests run here
+
+        Remote Namespace (_netns):
+            psp_dev_peer (2001:db8::2)
+            ★ psp_responder runs here
+
+    BPF programs:
+        - psp_dev_local: nk_forward.bpf.o (from parent class) uses bpf_redirect_peer()
+          to redirect directly to nk_guest (the peer of nk_host)
+        - nk_host: nk_redirect.bpf.o uses bpf_redirect_neigh() to redirect to
+          psp_dev_local with correct L2 headers via neighbor lookup
+    """
+
+    # Use a different prefix for netkit guest to avoid conflict with nsim prefix
+    nk_v6_pfx = "2001:db9::"
+
+    def __init__(self, src_path, **kwargs):
+        # Set LOCAL_PREFIX_V6 to a DIFFERENT prefix than nsim to avoid BPF
+        # redirecting psp_responder traffic. The BPF only redirects traffic
+        # matching LOCAL_PREFIX_V6, so nsim traffic (2001:db8::) won't be affected.
+        if "LOCAL_PREFIX_V6" not in os.environ:
+            os.environ["LOCAL_PREFIX_V6"] = self.nk_v6_pfx
+        super().__init__(src_path, **kwargs)
+
+        # Track nk_host BPF attachment for cleanup
+        self._nk_host_tc_attached = False
+        self._nk_host_bpf_prog_pref = None
+
+        # Set PSP device attributes based on whether we're using nsim or real NIC
+        if self._ns is not None:
+            # nsim case: PSP device is nsim_local (in host namespace)
+            self.psp_dev = self._ns.nsims[0].dev
+            self.psp_ifname = self.psp_dev['ifname']
+            self.psp_ifindex = self.psp_dev['ifindex']
+
+            # PSP peer device is nsim_peer (in _netns, where psp_responder runs)
+            self.psp_dev_peer = self._ns_peer.nsims[0].dev
+            self.psp_dev_peer_ifname = self.psp_dev_peer['ifname']
+            self.psp_dev_peer_ifindex = self.psp_dev_peer['ifindex']
+        else:
+            # Real NIC case: PSP device is the local interface
+            self.psp_dev = self.dev
+            self.psp_ifname = self.ifname
+            self.psp_ifindex = self.ifindex
+
+            # PSP peer device is the remote interface
+            self.psp_dev_peer = self.remote_dev
+            self.psp_dev_peer_ifname = self.remote_ifname
+            self.psp_dev_peer_ifindex = self.remote_ifindex
+
+        # Get nsid for the guest namespace (netns) where nk_guest is
+        self.psp_dev_peer_nsid = self._get_nsid(self.netns.name)
+
+        # Attach TX BPF on nk_host: redirect 2001:db8::/64 → nsim_local
+        self._attach_tx_bpf()
+
+        # Add routes for cross-namespace connectivity
+        self._setup_routes()
+
+    def _get_nsid(self, ns_name):
+        """Get the nsid for a namespace."""
+        nsid = 0
+        try:
+            list_result = cmd("ip netns list-id", shell=True)
+            list_output = list_result.stdout if hasattr(list_result, 'stdout') else str(list_result)
+
+            for line in list_output.split('\n'):
+                if ns_name in line and 'nsid' in line:
+                    parts = line.split()
+                    if 'nsid' in parts:
+                        nsid_idx = parts.index('nsid')
+                        if nsid_idx + 1 < len(parts):
+                            try:
+                                nsid = int(parts[nsid_idx + 1])
+                                break
+                            except ValueError:
+                                pass
+        except Exception as e:
+            raise KsftSkipEx(f"Failed to query nsid: {e}")
+        return nsid
+
+    def _setup_routes(self):
+        """
+        Set up routes for cross-namespace connectivity.
+
+        Traffic flows:
+        1. remote (_netns) -> nk_guest (netns):
+           nsim_peer -> nsim_local -> BPF redirect -> nk_host -> nk_guest
+           Needs: route in _netns to 2001:db9::/64 via nsim_local (2001:db8::1)
+
+        2. nk_guest (netns) -> remote (_netns):
+           nk_guest -> nk_host -> nsim_local -> nsim_peer
+           Needs: route in netns to 2001:db8::/64 via nk_host (default route already set)
+        """
+        # In _netns (remote namespace): add route to nk_guest prefix via nsim_local
+        # nsim_peer can reach nsim_local via the link, then traffic goes through BPF
+        ip(f"-6 route add {self.nk_v6_pfx}/64 via {self.nsim_v6_pfx}1 dev {self.psp_dev_peer_ifname}",
+           ns=self._netns)
+
+        # In netns (guest namespace): add route to nsim_peer prefix
+        # nk_guest default route goes to nk_host, but we need explicit route to 2001:db8::/64
+        ip(f"-6 route add {self.nsim_v6_pfx}/64 via fe80::1 dev {self._nk_guest_ifname}",
+           ns=self.netns)
+
+    def _attach_tx_bpf(self):
+        """
+        Attach BPF program on nk_host ingress to redirect TX traffic.
+
+        Packets from nk_guest destined for nsim network (2001:db8::/64) arrive
+        at nk_host via the netkit pair. This BPF program redirects them to
+        nsim_local so they can reach nsim_peer.
+        """
+        bpf_obj = self.test_dir / "nk_redirect.bpf.o"
+        if not bpf_obj.exists():
+            raise KsftSkipEx("BPF prog nk_redirect.bpf.o not found")
+
+        # Add clsact qdisc to nk_host
+        cmd(f"tc qdisc add dev {self._nk_host_ifname} clsact")
+
+        # Attach BPF on nk_host ingress
+        cmd(f"tc filter add dev {self._nk_host_ifname} ingress bpf obj {bpf_obj} sec tc/ingress direct-action")
+        self._nk_host_tc_attached = True
+
+        # Get the BPF program info
+        tc_info = cmd(f"tc filter show dev {self._nk_host_ifname} ingress").stdout
+        match = re.search(r'pref (\d+).*nk_redirect\.bpf.*id (\d+)', tc_info)
+        if not match:
+            raise Exception("Failed to get TX BPF prog ID")
+        self._nk_host_bpf_prog_pref = int(match.group(1))
+        nk_host_bpf_prog_id = int(match.group(2))
+
+        # Find and update the .bss map with nsim prefix and nsim_local ifindex
+        prog_info = bpftool(f"prog show id {nk_host_bpf_prog_id}", json=True)
+        map_ids = prog_info.get("map_ids", [])
+
+        bss_map_id = None
+        for map_id in map_ids:
+            map_info = bpftool(f"map show id {map_id}", json=True)
+            if map_info.get("name").endswith("bss"):
+                bss_map_id = map_id
+
+        if bss_map_id is None:
+            raise Exception("Failed to find TX BPF .bss map")
+
+        # Update map: match nsim prefix (2001:db8::), redirect to nsim_local ifindex
+        # BPF .bss layout (compiler reorders): ipv6_prefix (16 bytes) + redirect_ifindex (4 bytes)
+        ipv6_addr = ipaddress.IPv6Address(self.nsim_v6_pfx)
+        ipv6_bytes = ipv6_addr.packed
+        ifindex_bytes = self.psp_ifindex.to_bytes(4, byteorder='little')
+        value = ipv6_bytes + ifindex_bytes  # ipv6 first, then ifindex
+        value_hex = ' '.join(f'{b:02x}' for b in value)
+        bpftool(f"map update id {bss_map_id} key hex 00 00 00 00 value hex {value_hex}")
+
+    def __del__(self):
+        # Clean up nk_host BPF attachment
+        if hasattr(self, '_nk_host_tc_attached') and self._nk_host_tc_attached:
+            cmd(f"tc filter del dev {self._nk_host_ifname} ingress pref {self._nk_host_bpf_prog_pref}", fail=False)
+            self._nk_host_tc_attached = False
+
+        super().__del__()
diff --git a/tools/testing/selftests/drivers/net/psp.py b/tools/testing/selftests/drivers/net/psp.py
index 864d9fce1094..f6337028f41c 100755
--- a/tools/testing/selftests/drivers/net/psp.py
+++ b/tools/testing/selftests/drivers/net/psp.py
@@ -14,9 +14,11 @@ from lib.py import defer
 from lib.py import ksft_run, ksft_exit, ksft_pr
 from lib.py import ksft_true, ksft_eq, ksft_ne, ksft_gt, ksft_raises
 from lib.py import ksft_not_none
-from lib.py import KsftSkipEx
-from lib.py import NetDrvEpEnv, PSPFamily, NlError
-from lib.py import bkg, rand_port, wait_port_listen
+from lib.py import ksft_exit, ksft_pr, ksft_run, KsftSkipEx, ksft_true, KsftFailEx
+from lib.py import NetDrvEpEnv, NetDrvContEnvPsp, PSPFamily, NlError
+from lib.py import NetNSEnter
+from lib.py import bkg, cmd, rand_port, wait_port_listen, NetNSEnter
+from lib.py import ip
 
 
 def _get_outq(s):
@@ -117,11 +119,12 @@ def _get_stat(cfg, key):
 # Test case boiler plate
 #
 
-def _init_psp_dev(cfg):
+def _init_psp_dev(cfg, use_psp_ifindex=False):
     if not hasattr(cfg, 'psp_dev_id'):
         # Figure out which local device we are testing against
+        target_ifindex = cfg.psp_ifindex if use_psp_ifindex else cfg.ifindex
         for dev in cfg.pspnl.dev_get({}, dump=True):
-            if dev['ifindex'] == cfg.ifindex:
+            if dev['ifindex'] == target_ifindex:
                 cfg.psp_info = dev
                 cfg.psp_dev_id = cfg.psp_info['id']
                 break
@@ -394,6 +397,387 @@ def _data_basic_send(cfg, version, ipver):
     _close_psp_conn(cfg, s)
 
 
+def _data_basic_send_netkit_psp_assoc(cfg, version, ipver):
+    """
+    Test basic data send with netkit interface associated with PSP dev.
+    """
+
+    _init_psp_dev(cfg, True)
+    psp_dev_id_for_assoc = cfg.psp_dev_id
+
+    # Associate PSP device with nk_guest interface (in guest namespace)
+    nk_guest_dev = ip(f"link show dev {cfg._nk_guest_ifname}", json=True, ns=cfg.netns)[0]
+    nk_guest_ifindex = nk_guest_dev['ifindex']
+
+    cfg.pspnl.dev_assoc({'id': psp_dev_id_for_assoc, 'ifindex': nk_guest_ifindex, 'nsid': cfg.psp_dev_peer_nsid})
+
+    # Test connectivity in both directions before PSP operations
+    remote_addr = cfg.remote_addr_v["6"]  # nsim_peer address (2001:db8::2)
+    nk_guest_addr = cfg.nk_guest_ipv6     # nk_guest address (2001:db9::2:2)
+
+    # Check if assoc-list contains nk_guest
+    dev_info = cfg.pspnl.dev_get({'id': psp_dev_id_for_assoc})
+
+    if 'assoc-list' in dev_info:
+        found = False
+        for assoc in dev_info['assoc-list']:
+            if assoc['ifindex'] == nk_guest_ifindex and assoc['nsid'] == cfg.psp_dev_peer_nsid:
+                found = True
+                break
+        ksft_true(found, "Associated device not found in dev_get() response")
+    else:
+        raise RuntimeError("No assoc-list in dev_get() response after association")
+
+    # Enter guest namespace (netns) to run PSP test
+    with NetNSEnter(cfg.netns.name):
+        cfg.pspnl = PSPFamily()
+
+        s = _make_psp_conn(cfg, version, ipver)
+
+        rx_assoc = cfg.pspnl.rx_assoc({"version": version,
+                                       "dev-id": cfg.psp_dev_id,
+                                       "sock-fd": s.fileno()})
+        rx = rx_assoc['rx-key']
+        tx = _spi_xchg(s, rx)
+
+        cfg.pspnl.tx_assoc({"dev-id": cfg.psp_dev_id,
+                            "version": version,
+                            "tx-key": tx,
+                            "sock-fd": s.fileno()})
+
+        data_len = _send_careful(cfg, s, 100)
+        _check_data_rx(cfg, data_len)
+        _close_psp_conn(cfg, s)
+
+    # Clean up - back in host namespace
+    cfg.pspnl = PSPFamily()
+    cfg.pspnl.dev_disassoc({'id': psp_dev_id_for_assoc, 'ifindex': nk_guest_ifindex, 'nsid': cfg.psp_dev_peer_nsid})
+
+    del cfg.psp_dev_id
+    del cfg.psp_info
+
+
+def _key_rotation_notify_multi_ns_netkit(cfg, version, ipver):
+    """ Test key rotation notifications across multiple namespaces using netkit """
+    import threading
+
+    _init_psp_dev(cfg, True)
+    psp_dev_id_for_assoc = cfg.psp_dev_id
+
+    # Associate PSP device with nk_guest interface (in guest namespace)
+    nk_guest_dev = ip(f"link show dev {cfg._nk_guest_ifname}", json=True, ns=cfg.netns)[0]
+    nk_guest_ifindex = nk_guest_dev['ifindex']
+
+    cfg.pspnl.dev_assoc({'id': psp_dev_id_for_assoc, 'ifindex': nk_guest_ifindex, 'nsid': cfg.psp_dev_peer_nsid})
+
+    results = {'main_ns': None, 'peer_ns': None}
+
+    def listen_in_namespace(ns_name, result_key):
+        """Listen for key rotation notifications in a namespace"""
+        try:
+            if ns_name:
+                ctx = NetNSEnter(ns_name)
+            else:
+                from contextlib import nullcontext
+                ctx = nullcontext()
+
+            with ctx:
+                pspnl = PSPFamily()
+                pspnl.ntf_subscribe('use')
+
+                collected_notifications = []
+                for i in range(100):
+                    pspnl.check_ntf()
+
+                    try:
+                        while True:
+                            msg = pspnl.async_msg_queue.get_nowait()
+                            collected_notifications.append(msg['msg'])
+                    except:
+                        pass
+
+                    if collected_notifications:
+                        results[result_key] = collected_notifications
+                        break
+
+                    time.sleep(0.1)
+                else:
+                    # Timeout - no notification received
+                    results[result_key] = []
+        except Exception as e:
+            results[result_key] = {'error': str(e)}
+
+    # Create listener threads
+    # Main namespace listener
+    main_ns_thread = threading.Thread(
+        target=listen_in_namespace,
+        args=(None, 'main_ns'),
+        name='main_ns_listener'
+    )
+    # Guest namespace listener (netkit peer)
+    guest_ns_thread = threading.Thread(
+        target=listen_in_namespace,
+        args=(cfg.netns.name, 'peer_ns'),
+        name='guest_ns_listener'
+    )
+
+    main_ns_thread.start()
+    guest_ns_thread.start()
+
+    time.sleep(0.5)
+
+    # Trigger key rotation on the PSP device
+    cfg.pspnl.key_rotate({"id": psp_dev_id_for_assoc})
+
+    main_ns_thread.join(timeout=15)
+    guest_ns_thread.join(timeout=15)
+
+    # Verify notifications were received in both namespaces
+    ksft_not_none(results['main_ns'], "No result from main namespace listener")
+    ksft_not_none(results['peer_ns'], "No result from guest namespace listener")
+
+    # Check for errors in listeners
+    if isinstance(results['main_ns'], dict) and 'error' in results['main_ns']:
+        raise RuntimeError(f"Main NS listener error: {results['main_ns']['error']}")
+    if isinstance(results['peer_ns'], dict) and 'error' in results['peer_ns']:
+        raise RuntimeError(f"Guest NS listener error: {results['peer_ns']['error']}")
+
+    # Verify that both namespaces received at least one notification
+    ksft_gt(len(results['main_ns']), 0, "No key rotation notification received in main namespace")
+    ksft_gt(len(results['peer_ns']), 0, "No key rotation notification received in guest namespace")
+
+    # Verify the notification contains the expected dev_id
+    main_ns_has_dev = any(
+        ntf.get('id') == psp_dev_id_for_assoc
+        for ntf in results['main_ns']
+    )
+    guest_ns_has_dev = any(
+        ntf.get('id') == psp_dev_id_for_assoc
+        for ntf in results['peer_ns']
+    )
+
+    ksft_true(main_ns_has_dev, "Key rotation notification for correct device not found in main namespace")
+    ksft_true(guest_ns_has_dev, "Key rotation notification for correct device not found in guest namespace")
+
+    # Clean up
+    cfg.pspnl.dev_disassoc({'id': psp_dev_id_for_assoc, 'ifindex': nk_guest_ifindex, 'nsid': cfg.psp_dev_peer_nsid})
+    del cfg.psp_dev_id
+    del cfg.psp_info
+
+
+def _dev_change_notify_multi_ns_netkit(cfg, version, ipver):
+    """ Test dev_change notifications across multiple namespaces using netkit """
+    import threading
+
+    _init_psp_dev(cfg, True)
+    psp_dev_id_for_assoc = cfg.psp_dev_id
+
+    # Associate PSP device with nk_guest interface (in guest namespace)
+    nk_guest_dev = ip(f"link show dev {cfg._nk_guest_ifname}", json=True, ns=cfg.netns)[0]
+    nk_guest_ifindex = nk_guest_dev['ifindex']
+
+    cfg.pspnl.dev_assoc({'id': psp_dev_id_for_assoc, 'ifindex': nk_guest_ifindex, 'nsid': cfg.psp_dev_peer_nsid})
+
+    results = {'main_ns': None, 'peer_ns': None}
+
+    def listen_in_namespace(ns_name, result_key):
+        """Listen for dev_change notifications in a namespace"""
+        try:
+            if ns_name:
+                ctx = NetNSEnter(ns_name)
+            else:
+                from contextlib import nullcontext
+                ctx = nullcontext()
+
+            with ctx:
+                pspnl = PSPFamily()
+                pspnl.ntf_subscribe('mgmt')
+
+                collected_notifications = []
+                for i in range(100):
+                    pspnl.check_ntf()
+
+                    try:
+                        while True:
+                            msg = pspnl.async_msg_queue.get_nowait()
+                            collected_notifications.append(msg['msg'])
+                    except:
+                        pass
+
+                    if collected_notifications:
+                        results[result_key] = collected_notifications
+                        break
+
+                    time.sleep(0.1)
+                else:
+                    # Timeout - no notification received
+                    results[result_key] = []
+        except Exception as e:
+            results[result_key] = {'error': str(e)}
+
+    # Create listener threads
+    # Main namespace listener
+    main_ns_thread = threading.Thread(
+        target=listen_in_namespace,
+        args=(None, 'main_ns'),
+        name='main_ns_listener'
+    )
+    # Guest namespace listener (netkit peer)
+    guest_ns_thread = threading.Thread(
+        target=listen_in_namespace,
+        args=(cfg.netns.name, 'peer_ns'),
+        name='guest_ns_listener'
+    )
+
+    main_ns_thread.start()
+    guest_ns_thread.start()
+
+    time.sleep(1.0)  # Give threads time to subscribe
+
+    # Trigger dev_change by calling dev_set (notification is always sent)
+    cfg.pspnl.dev_set({'id': psp_dev_id_for_assoc, 'psp-versions-ena': cfg.psp_info['psp-versions-cap']})
+
+    main_ns_thread.join(timeout=15)
+    guest_ns_thread.join(timeout=15)
+
+    # Verify notifications were received in both namespaces
+    ksft_not_none(results['main_ns'], "No result from main namespace listener")
+    ksft_not_none(results['peer_ns'], "No result from guest namespace listener")
+
+    # Check for errors in listeners
+    if isinstance(results['main_ns'], dict) and 'error' in results['main_ns']:
+        raise RuntimeError(f"Main NS listener error: {results['main_ns']['error']}")
+    if isinstance(results['peer_ns'], dict) and 'error' in results['peer_ns']:
+        raise RuntimeError(f"Guest NS listener error: {results['peer_ns']['error']}")
+
+    # Verify that both namespaces received at least one notification
+    ksft_gt(len(results['main_ns']), 0, "No dev_change notification received in main namespace")
+    ksft_gt(len(results['peer_ns']), 0, "No dev_change notification received in guest namespace")
+
+    # Verify the notification contains the expected dev_id
+    main_ns_has_dev = any(
+        ntf.get('id') == psp_dev_id_for_assoc
+        for ntf in results['main_ns']
+    )
+    guest_ns_has_dev = any(
+        ntf.get('id') == psp_dev_id_for_assoc
+        for ntf in results['peer_ns']
+    )
+
+    ksft_true(main_ns_has_dev, "Dev_change notification for correct device not found in main namespace")
+    ksft_true(guest_ns_has_dev, "Dev_change notification for correct device not found in guest namespace")
+
+    # Clean up
+    cfg.pspnl.dev_disassoc({'id': psp_dev_id_for_assoc, 'ifindex': nk_guest_ifindex, 'nsid': cfg.psp_dev_peer_nsid})
+    del cfg.psp_dev_id
+    del cfg.psp_info
+
+
+def _psp_dev_get_check_netkit_psp_assoc(cfg, version, ipver):
+    """ Check psp dev-get output with netkit interface associated with PSP dev """
+
+    _init_psp_dev(cfg, True)
+    psp_dev_id_for_assoc = cfg.psp_dev_id
+
+    # Associate PSP device with nk_guest interface (in guest namespace)
+    nk_guest_dev = ip(f"link show dev {cfg._nk_guest_ifname}", json=True, ns=cfg.netns)[0]
+    nk_guest_ifindex = nk_guest_dev['ifindex']
+
+    cfg.pspnl.dev_assoc({'id': psp_dev_id_for_assoc, 'ifindex': nk_guest_ifindex, 'nsid': cfg.psp_dev_peer_nsid})
+
+    # Check 1: In default netns, verify dev-get has correct ifindex and assoc-list
+    dev_info = cfg.pspnl.dev_get({'id': psp_dev_id_for_assoc})
+
+    # Verify the PSP device has the correct ifindex
+    ksft_eq(dev_info['ifindex'], cfg.psp_ifindex)
+
+    # Verify assoc-list exists and contains the associated nk_guest with correct ifindex and nsid
+    ksft_true('assoc-list' in dev_info, "No assoc-list in dev_get() response after association")
+    found = False
+    for assoc in dev_info['assoc-list']:
+        if assoc['ifindex'] == nk_guest_ifindex and assoc['nsid'] == cfg.psp_dev_peer_nsid:
+            found = True
+            break
+    ksft_true(found, "Associated device not found in assoc-list with correct ifindex and nsid")
+
+    # Check 2: In guest netns, verify dev-get has assoc-list with nk_guest device
+    with NetNSEnter(cfg.netns.name):
+        peer_pspnl = PSPFamily()
+
+        # Dump all devices in the guest namespace
+        peer_devices = peer_pspnl.dev_get({}, dump=True)
+
+        # Find the device with by-association flag
+        peer_dev = None
+        for dev in peer_devices:
+            if dev.get('by-association'):
+                peer_dev = dev
+                break
+
+        ksft_not_none(peer_dev, "No PSP device found with by-association flag in guest netns")
+
+        # Verify assoc-list contains the nk_guest device
+        ksft_true('assoc-list' in peer_dev and len(peer_dev['assoc-list']) > 0,
+                  "Guest device should have assoc-list with local devices")
+
+        # Verify the assoc-list contains nk_guest ifindex with nsid=-1 (same namespace)
+        found = False
+        for assoc in peer_dev['assoc-list']:
+            if assoc['ifindex'] == nk_guest_ifindex:
+                ksft_eq(assoc['nsid'], -1,
+                        "nsid should be -1 (NETNSA_NSID_NOT_ASSIGNED) for same-namespace device")
+                found = True
+                break
+        ksft_true(found, "nk_guest ifindex not found in assoc-list")
+
+    # Clean up
+    cfg.pspnl.dev_disassoc({'id': psp_dev_id_for_assoc, 'ifindex': nk_guest_ifindex, 'nsid': cfg.psp_dev_peer_nsid})
+
+    del cfg.psp_dev_id
+    del cfg.psp_info
+
+
+def _psp_dev_assoc_cleanup_on_netkit_del(cfg):
+    """ Test that assoc-list is cleared when associated netkit interface is deleted """
+    import subprocess
+
+    _init_psp_dev(cfg, True)
+    psp_dev_id_for_assoc = cfg.psp_dev_id
+
+    # Associate PSP device with nk_guest interface (in guest namespace)
+    nk_guest_dev = ip(f"link show dev {cfg._nk_guest_ifname}", json=True, ns=cfg.netns)[0]
+    nk_guest_ifindex = nk_guest_dev['ifindex']
+
+    cfg.pspnl.dev_assoc({'id': psp_dev_id_for_assoc, 'ifindex': nk_guest_ifindex, 'nsid': cfg.psp_dev_peer_nsid})
+
+    # Verify assoc-list exists in default netns
+    dev_info = cfg.pspnl.dev_get({'id': psp_dev_id_for_assoc})
+    ksft_true('assoc-list' in dev_info, "No assoc-list after association")
+    found = False
+    for assoc in dev_info['assoc-list']:
+        if assoc['ifindex'] == nk_guest_ifindex and assoc['nsid'] == cfg.psp_dev_peer_nsid:
+            found = True
+            break
+    ksft_true(found, "Associated device not found in assoc-list")
+
+    # Delete the netkit interface in the guest namespace
+    ns_name = cfg.netns.name
+    subprocess.run(f"ip netns exec {ns_name} ip link del {cfg._nk_guest_ifname}", shell=True, check=True)
+
+    # Mark netkit as already deleted so cleanup won't try to delete it again
+    # (deleting nk_guest also removes nk_host since they're a pair)
+    cfg._nk_host_ifname = None
+    cfg._nk_guest_ifname = None
+
+    # Verify assoc-list is gone in default netns after netkit deletion
+    dev_info = cfg.pspnl.dev_get({'id': psp_dev_id_for_assoc})
+    ksft_true('assoc-list' not in dev_info or len(dev_info['assoc-list']) == 0,
+              "assoc-list should be empty after netkit deletion")
+
+    del cfg.psp_dev_id
+    del cfg.psp_info
+
+
 def __bad_xfer_do(cfg, s, tx, version='hdr0-aes-gcm-128'):
     # Make sure we accept the ACK for the SPI before we seal with the bad assoc
     _check_data_outq(s, 0)
@@ -637,6 +1021,74 @@ def main() -> None:
                 ksft_pr("STDOUT:\n#  " + srv.stdout.strip().replace("\n", "\n#  "))
             if srv and srv.stderr:
                 ksft_pr("STDERR:\n#  " + srv.stderr.strip().replace("\n", "\n#  "))
+
+    with NetDrvContEnvPsp(__file__) as cfg:
+        cfg.pspnl = PSPFamily()
+
+        # Set up responder and communication sock
+        # For NetDrvContEnvPsp: psp_responder runs in _netns (remote namespace with nsim_peer)
+        responder = cfg.remote.deploy("psp_responder")
+
+        cfg.comm_port = rand_port()
+        srv = None
+        try:
+            with bkg(responder + f" -p {cfg.comm_port}", host=cfg.remote,
+                     exit_wait=True) as srv:
+                wait_port_listen(cfg.comm_port, host=cfg.remote)
+
+                cfg.comm_sock = socket.create_connection((cfg.remote_addr,
+                                                          cfg.comm_port),
+                                                         timeout=1)
+
+                cases = [
+                    psp_ip_ver_test_builder(
+                        "data_basic_send_netkit_psp_assoc",
+                        _data_basic_send_netkit_psp_assoc, version, "6"
+                    )
+                    for version in range(0, 4)
+                ]
+                cases += [
+                    psp_ip_ver_test_builder(
+                        "key_rotation_notify_multi_ns_netkit",
+                        _key_rotation_notify_multi_ns_netkit, version, "6"
+                    )
+                    for version in range(0, 4)
+                ]
+                cases += [
+                    psp_ip_ver_test_builder(
+                        "dev_change_notify_multi_ns_netkit",
+                        _dev_change_notify_multi_ns_netkit, version, "6"
+                    )
+                    for version in range(0, 4)
+                ]
+                cases += [
+                    psp_ip_ver_test_builder(
+                        "psp_dev_get_check_netkit_psp_assoc",
+                        _psp_dev_get_check_netkit_psp_assoc, version, "6"
+                    )
+                    for version in range(0, 4)
+                ]
+                # Run netkit deletion test only once at the end since it destroys the netkit
+                def psp_dev_assoc_cleanup_on_netkit_del_test(cfg):
+                    _psp_dev_assoc_cleanup_on_netkit_del(cfg)
+                psp_dev_assoc_cleanup_on_netkit_del_test.__name__ = "psp_dev_assoc_cleanup_on_netkit_del"
+                cases.append(psp_dev_assoc_cleanup_on_netkit_del_test)
+
+                ksft_run(cases=cases, globs=globals(),
+                         case_pfx={},
+                         args=(cfg, ))
+
+                cfg.comm_sock.send(b"exit\0")
+                cfg.comm_sock.close()
+        finally:
+            if srv and (srv.stdout or srv.stderr):
+                ksft_pr("")
+                ksft_pr(f"Responder logs ({srv.ret}):")
+            if srv and srv.stdout:
+                ksft_pr("STDOUT:\n#  " + srv.stdout.strip().replace("\n", "\n#  "))
+            if srv and srv.stderr:
+                ksft_pr("STDERR:\n#  " + srv.stderr.strip().replace("\n", "\n#  "))
+
     ksft_exit()
 
 
-- 
2.47.3