diff --git a/tests/hwsim/test_ap_pmf.py b/tests/hwsim/test_ap_pmf.py
index 15777b21a..446172919 100644
--- a/tests/hwsim/test_ap_pmf.py
+++ b/tests/hwsim/test_ap_pmf.py
@@ -12,6 +12,7 @@ logger = logging.getLogger()
 import hwsim_utils
 import hostapd
 from wlantest import Wlantest
+from wpasupplicant import WpaSupplicant
 
 def test_ap_pmf_required(dev, apdev):
     """WPA2-PSK AP with PMF required"""
@@ -114,3 +115,72 @@ def test_ap_pmf_negative(dev, apdev):
     except Exception, e:
         logger.debug("Ignore expected exception: " + str(e))
     wt.require_ap_no_pmf(apdev[0]['bssid'])
+
+def test_ap_pmf_assoc_comeback(dev, apdev):
+    """WPA2-PSK AP with PMF association comeback"""
+    ssid = "assoc-comeback"
+    wt = Wlantest()
+    wt.flush()
+    wt.add_passphrase("12345678")
+    params = hostapd.wpa2_params(ssid=ssid, passphrase="12345678")
+    params["wpa_key_mgmt"] = "WPA-PSK-SHA256";
+    params["ieee80211w"] = "2";
+    hapd = hostapd.add_ap(apdev[0]['ifname'], params)
+    dev[0].connect(ssid, psk="12345678", ieee80211w="1",
+                   key_mgmt="WPA-PSK WPA-PSK-SHA256", proto="WPA2",
+                   scan_freq="2412")
+    hapd.set("ext_mgmt_frame_handling", "1")
+    dev[0].request("DISCONNECT")
+    ev = dev[0].wait_event(["CTRL-EVENT-DISCONNECTED"])
+    if ev is None:
+        raise Exception("Timeout on disconnection")
+    hapd.set("ext_mgmt_frame_handling", "0")
+    dev[0].request("REASSOCIATE")
+    ev = dev[0].wait_event(["CTRL-EVENT-CONNECTED"])
+    if ev is None:
+        raise Exception("Timeout on re-connection")
+    if wt.get_sta_counter("assocresp_comeback", apdev[0]['bssid'],
+                          dev[0].p2p_interface_addr()) < 1:
+        raise Exception("AP did not use association comeback request")
+
+def test_ap_pmf_sta_sa_query(dev, apdev):
+    """WPA2-PSK AP with station using SA Query"""
+    ssid = "assoc-comeback"
+    addr = dev[0].p2p_dev_addr()
+    wt = Wlantest()
+    wt.flush()
+    wt.add_passphrase("12345678")
+
+    wpas = WpaSupplicant(global_iface='/tmp/wpas-wlan5')
+    wpas.interface_add("wlan5", drv_params="use_monitor=1")
+    id = wpas.add_network()
+    wpas.set_network(id, "mode", "2")
+    wpas.set_network_quoted(id, "ssid", ssid)
+    wpas.set_network(id, "proto", "WPA2")
+    wpas.set_network(id, "key_mgmt", "WPA-PSK-SHA256")
+    wpas.set_network(id, "ieee80211w", "2")
+    wpas.set_network_quoted(id, "psk", "12345678")
+    wpas.set_network(id, "pairwise", "CCMP")
+    wpas.set_network(id, "group", "CCMP")
+    wpas.set_network(id, "frequency", "2412")
+    wpas.connect_network(id)
+    bssid = wpas.p2p_dev_addr()
+
+    dev[0].connect(ssid, psk="12345678", ieee80211w="1",
+                   key_mgmt="WPA-PSK WPA-PSK-SHA256", proto="WPA2",
+                   scan_freq="2412")
+    wpas.request("DEAUTHENTICATE " + addr + " test=0")
+    wpas.request("DISASSOCIATE " + addr + " test=0")
+    ev = dev[0].wait_event(["CTRL-EVENT-DISCONNECTED"], timeout=1)
+    if ev is not None:
+        raise Exception("Unexpected disconnection")
+
+    wpas.request("DEAUTHENTICATE " + addr + " reason=6 test=0")
+    wpas.request("DISASSOCIATE " + addr + " reason=7 test=0")
+    ev = dev[0].wait_event(["CTRL-EVENT-DISCONNECTED"], timeout=1)
+    if ev is not None:
+        raise Exception("Unexpected disconnection")
+    if wt.get_sta_counter("valid_saqueryreq_tx", bssid, addr) < 1:
+        raise Exception("STA did not send SA Query")
+    if wt.get_sta_counter("valid_saqueryresp_rx", bssid, addr) < 1:
+        raise Exception("AP did not reply to SA Query")