[connectivity_test_utils] Utils methods for socket keepalive tests

CHERRY PICKED CHANGES FROM AOSP

Bug: 123774509
Test: Verified the changes
Change-Id: I582cc33403d636638572ea2d919c5a9a3b5c286f
Merged-In: Ie15c80db415d2af808b55674a1402983933ddaf2
diff --git a/acts/framework/acts/test_utils/net/connectivity_test_utils.py b/acts/framework/acts/test_utils/net/connectivity_test_utils.py
index 4b7668c..153630f 100644
--- a/acts/framework/acts/test_utils/net/connectivity_test_utils.py
+++ b/acts/framework/acts/test_utils/net/connectivity_test_utils.py
@@ -15,54 +15,95 @@
 
 from acts import asserts
 from acts.test_utils.net import connectivity_const as cconst
+from queue import Empty
 
-def start_natt_keepalive(ad, src_ip, src_port, dst_ip, interval = 10):
-    """ Start NAT-T keep alive on dut """
+def _listen_for_keepalive_event(ad, key, msg, ka_event):
+    """Listen for keepalive event and return status
 
-    ad.log.info("Starting NATT Keepalive")
-    status = None
-
-    key = ad.droid.connectivityStartNattKeepalive(
-        interval, src_ip, src_port, dst_ip)
-
-    ad.droid.connectivityNattKeepaliveStartListeningForEvent(key, "Started")
+    Args:
+        ad: DUT object
+        key: keepalive key
+        msg: Error message
+        event: Keepalive event type
+    """
+    ad.droid.socketKeepaliveStartListeningForEvent(key, ka_event)
     try:
-        event = ad.ed.pop_event("PacketKeepaliveCallback")
-        status = event["data"]["packetKeepaliveEvent"]
+        event = ad.ed.pop_event("SocketKeepaliveCallback")
+        status = event["data"]["socketKeepaliveEvent"] == ka_event
     except Empty:
-        msg = "Failed to receive confirmation of starting natt keepalive"
         asserts.fail(msg)
     finally:
-        ad.droid.connectivityNattKeepaliveStopListeningForEvent(
-            key, "Started")
-
-    if status != "Started":
-        ad.log.error("Received keepalive status: %s" % status)
-        ad.droid.connectivityRemovePacketKeepaliveReceiverKey(key)
-        return None
-    return key
-
-def stop_natt_keepalive(ad, key):
-    """ Stop NAT-T keep alive on dut """
-
-    ad.log.info("Stopping NATT keepalive")
-    status = False
-    ad.droid.connectivityStopNattKeepalive(key)
-
-    ad.droid.connectivityNattKeepaliveStartListeningForEvent(key, "Stopped")
-    try:
-        event = ad.ed.pop_event("PacketKeepaliveCallback")
-        status = event["data"]["packetKeepaliveEvent"] == "Stopped"
-    except Empty:
-        msg = "Failed to receive confirmation of stopping natt keepalive"
-        asserts.fail(msg)
-    finally:
-        ad.droid.connectivityNattKeepaliveStopListeningForEvent(
-            key, "Stopped")
-
-    ad.droid.connectivityRemovePacketKeepaliveReceiverKey(key)
+        ad.droid.socketKeepaliveStopListeningForEvent(key, ka_event)
+    if ka_event != "Started":
+        ad.droid.removeSocketKeepaliveReceiverKey(key)
+    if status:
+        ad.log.info("'%s' keepalive event successful" % ka_event)
     return status
 
+def start_natt_socket_keepalive(ad, udp_encap, src_ip, dst_ip, interval = 10):
+    """Start NATT SocketKeepalive on DUT
+
+    Args:
+        ad: DUT object
+        udp_encap: udp_encap socket key
+        src_ip: IP addr of the client
+        dst_ip: IP addr of the keepalive server
+        interval: keepalive time interval
+    """
+    ad.log.info("Starting Natt Socket Keepalive")
+    key = ad.droid.startNattSocketKeepalive(udp_encap, src_ip, dst_ip, interval)
+    msg = "Failed to receive confirmation of starting natt socket keeaplive"
+    status = _listen_for_keepalive_event(ad, key, msg, "Started")
+    return key if status else None
+
+def start_tcp_socket_keepalive(ad, socket, time_interval = 10):
+    """Start TCP socket keepalive on DUT
+
+    Args:
+        ad: DUT object
+        socket: TCP socket key
+        time_interval: Keepalive time interval
+    """
+    ad.log.info("Starting TCP Socket Keepalive")
+    key = ad.droid.startTcpSocketKeepalive(socket, time_interval)
+    msg = "Failed to receive confirmation of starting tcp socket keeaplive"
+    status = _listen_for_keepalive_event(ad, key, msg, "Started")
+    return key if status else None
+
+def socket_keepalive_error(ad, key):
+    """Verify Error callback
+
+    Args:
+        ad: DUT object
+        key: Keepalive key
+    """
+    ad.log.info("Verify Error callback on keepalive: %s" % key)
+    msg = "Failed to receive confirmation of Error callback"
+    return _listen_for_keepalive_event(ad, key, msg, "Error")
+
+def socket_keepalive_data_received(ad, key):
+    """Verify OnDataReceived callback
+
+    Args:
+        ad: DUT object
+        key: Keepalive key
+    """
+    ad.log.info("Verify OnDataReceived callback on keepalive: %s" % key)
+    msg = "Failed to receive confirmation of OnDataReceived callback"
+    return _listen_for_keepalive_event(ad, key, msg, "OnDataReceived")
+
+def stop_socket_keepalive(ad, key):
+    """Stop SocketKeepalive on DUT
+
+    Args:
+        ad: DUT object
+        key: Keepalive key
+    """
+    ad.log.info("Stopping Socket keepalive: %s" % key)
+    ad.droid.stopSocketKeepalive(key)
+    msg = "Failed to receive confirmation of stopping socket keepalive"
+    return _listen_for_keepalive_event(ad, key, msg, "Stopped")
+
 def set_private_dns(ad, dns_mode, hostname=None):
     """ Set private DNS mode on dut """
     if dns_mode == cconst.PRIVATE_DNS_MODE_OFF: