Merge "Framework side implementation of carrier roaming satellite metrics." into 24D1-dev
diff --git a/src/java/com/android/internal/telephony/data/DataNetwork.java b/src/java/com/android/internal/telephony/data/DataNetwork.java
index bd98403..7ca504b 100644
--- a/src/java/com/android/internal/telephony/data/DataNetwork.java
+++ b/src/java/com/android/internal/telephony/data/DataNetwork.java
@@ -124,6 +124,7 @@
 import java.util.Comparator;
 import java.util.HashSet;
 import java.util.List;
+import java.util.ListIterator;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
@@ -1001,23 +1002,36 @@
                 && transport == AccessNetworkConstants.TRANSPORT_TYPE_WWAN;
         mDataAllowedReason = dataAllowedReason;
         dataProfile.setLastSetupTimestamp(SystemClock.elapsedRealtime());
-        mAttachedNetworkRequestList.addAll(networkRequestList);
         for (int transportType : mAccessNetworksManager.getAvailableTransports()) {
             mCid.put(transportType, INVALID_CID);
         }
         mTelephonyDisplayInfo = mPhone.getDisplayInfoController().getTelephonyDisplayInfo();
         mTcpBufferSizes = mDataConfigManager.getTcpConfigString(mTelephonyDisplayInfo);
 
-        for (TelephonyNetworkRequest networkRequest : networkRequestList) {
-            networkRequest.setAttachedNetwork(DataNetwork.this);
-            networkRequest.setState(TelephonyNetworkRequest.REQUEST_STATE_SATISFIED);
-        }
-
+        // network capabilities infer connectivity transport and MMTEL from the requested
+        // capabilities.
+        // TODO: Ideally we shouldn't infer network capabilities base on the requested capabilities,
+        // but currently there are 2 hacks associated with getForcedCellularTransportCapabilities
+        // and IMS service requesting IMS|MMTEL that need to support. When we stop supporting these
+        // cases, we shouldn't consider the requests when determining the network capabilities.
+        mAttachedNetworkRequestList.addAll(networkRequestList);
         // Update the capabilities in the constructor is to make sure the data network has initial
         // capability immediately after created. Doing this connecting state creates the window that
         // DataNetworkController might check if existing data network's capability can satisfy the
         // next network request within this window.
         updateNetworkCapabilities();
+
+        // Remove the requests that can't use the initial capabilities
+        ListIterator<TelephonyNetworkRequest> iter = mAttachedNetworkRequestList.listIterator();
+        while (iter.hasNext()) {
+            TelephonyNetworkRequest request = iter.next();
+            if (request.canBeSatisfiedBy(mNetworkCapabilities)) {
+                request.setAttachedNetwork(DataNetwork.this);
+                request.setState(TelephonyNetworkRequest.REQUEST_STATE_SATISFIED);
+            } else {
+                iter.remove();
+            }
+        }
     }
 
     /**
diff --git a/src/java/com/android/internal/telephony/data/DataNetworkController.java b/src/java/com/android/internal/telephony/data/DataNetworkController.java
index 2725295..5a688ca 100644
--- a/src/java/com/android/internal/telephony/data/DataNetworkController.java
+++ b/src/java/com/android/internal/telephony/data/DataNetworkController.java
@@ -559,6 +559,16 @@
         }
 
         /**
+         * Print "capabilities - connectivity transport". e.g. INTERNET|NOT_RESTRICTED-SATELLITE
+         */
+        @NonNull
+        public String toStringSimplified() {
+            return size() > 0 ? DataUtils.networkCapabilitiesToString(get(0).getCapabilities())
+                    + "-" + DataUtils.connectivityTransportsToString(get(0).getTransportTypes())
+                    : "";
+        }
+
+        /**
          * Dump the network request list.
          *
          * @param pw print writer.
@@ -1838,8 +1848,7 @@
         log("Re-evaluating " + networkRequestLists.stream().mapToInt(List::size).sum()
                 + " unsatisfied network requests in " + networkRequestLists.size()
                 + " groups, " + networkRequestLists.stream().map(
-                        requestList -> DataUtils.networkCapabilitiesToString(
-                                requestList.get(0).getCapabilities()))
+                        NetworkRequestList::toStringSimplified)
                 .collect(Collectors.joining(", ")) + " due to " + reason);
 
         // Second, see if any existing network can satisfy those network requests.
@@ -2113,35 +2122,30 @@
      */
     private boolean canConnectivityTransportSatisfyNetworkRequest(
             @NonNull TelephonyNetworkRequest networkRequest, @TransportType int transport) {
+        // Check if this is a IWLAN network request.
+        if (transport == AccessNetworkConstants.TRANSPORT_TYPE_WLAN) {
+            // If the request would result in bringing up network on IWLAN, then no
+            // need to check if the device is using satellite network.
+            return true;
+        }
+
         // When the device is on satellite, only restricted network request can request network.
         if (mServiceState.isUsingNonTerrestrialNetwork()
-                && networkRequest.getNativeNetworkRequest().hasCapability(
+                && networkRequest.hasCapability(
                         NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)) {
             return false;
         }
 
         // If the network request does not specify cellular or satellite, then it can be
         // satisfied when the device is either on cellular ot satellite.
-        if (!networkRequest.getNativeNetworkRequest().hasTransport(
-                NetworkCapabilities.TRANSPORT_CELLULAR)
-                && !networkRequest.getNativeNetworkRequest().hasTransport(
-                        NetworkCapabilities.TRANSPORT_SATELLITE)) {
-            return true;
-        }
-
-        // Check if this is a IWLAN network request.
-        if (networkRequest.getNativeNetworkRequest().hasTransport(
-                NetworkCapabilities.TRANSPORT_CELLULAR)
-                && transport == AccessNetworkConstants.TRANSPORT_TYPE_WLAN) {
-            // If the cellular request would result in bringing up network on IWLAN, then no
-            // need to check if the device is using satellite network.
+        if (!networkRequest.hasTransport(NetworkCapabilities.TRANSPORT_CELLULAR)
+                && !networkRequest.hasTransport(NetworkCapabilities.TRANSPORT_SATELLITE)) {
             return true;
         }
 
         // As a short term solution, allowing some networks to be always marked as cellular
         // transport if certain capabilities are in the network request.
-        if (networkRequest.getNativeNetworkRequest().hasTransport(
-                NetworkCapabilities.TRANSPORT_CELLULAR) && Arrays.stream(
+        if (networkRequest.hasTransport(NetworkCapabilities.TRANSPORT_CELLULAR) && Arrays.stream(
                         networkRequest.getCapabilities())
                 .anyMatch(mDataConfigManager.getForcedCellularTransportCapabilities()::contains)) {
             return true;
@@ -2151,11 +2155,9 @@
         // the network is satellite, then the request must specify satellite transport and
         // restricted.
         return (mServiceState.isUsingNonTerrestrialNetwork()
-                && networkRequest.getNativeNetworkRequest().hasTransport(
-                        NetworkCapabilities.TRANSPORT_SATELLITE))
+                && networkRequest.hasTransport(NetworkCapabilities.TRANSPORT_SATELLITE))
                 || (!mServiceState.isUsingNonTerrestrialNetwork()
-                        && networkRequest.getNativeNetworkRequest().hasTransport(
-                        NetworkCapabilities.TRANSPORT_CELLULAR));
+                        && networkRequest.hasTransport(NetworkCapabilities.TRANSPORT_CELLULAR));
     }
 
     /**
@@ -3058,8 +3060,18 @@
             List<NetworkRequestList> groupRequestLists = getGroupedUnsatisfiedNetworkRequests();
             dataSetupRetryEntry.networkRequestList.stream()
                     .filter(request -> groupRequestLists.stream()
-                            .anyMatch(groupRequestList -> groupRequestList
-                                    .get(request.getCapabilities()) != null))
+                            .anyMatch(groupRequestList -> {
+                                // The unsatisfied request has all the requested capabilities.
+                                if (groupRequestList.get(request.getCapabilities()) == null) {
+                                    return false;
+                                }
+                                TelephonyNetworkRequest leading = groupRequestList.getFirst();
+                                // The unsatisfied request covers all the requested transports.
+                                return leading.getTransportTypes().length == 0
+                                        || request.getTransportTypes().length == 0
+                                        || Arrays.stream(request.getTransportTypes())
+                                        .allMatch(leading::hasTransport);
+                            }))
                     .forEach(requestList::add);
         }
         if (requestList.isEmpty()) {
diff --git a/src/java/com/android/internal/telephony/data/DataUtils.java b/src/java/com/android/internal/telephony/data/DataUtils.java
index cca6deb..6814516 100644
--- a/src/java/com/android/internal/telephony/data/DataUtils.java
+++ b/src/java/com/android/internal/telephony/data/DataUtils.java
@@ -27,6 +27,7 @@
 import android.telephony.AccessNetworkConstants.AccessNetworkType;
 import android.telephony.AccessNetworkConstants.RadioAccessNetworkType;
 import android.telephony.AccessNetworkConstants.TransportType;
+import android.telephony.Annotation.ConnectivityTransport;
 import android.telephony.Annotation.DataActivityType;
 import android.telephony.Annotation.NetCapability;
 import android.telephony.Annotation.NetworkType;
@@ -176,6 +177,44 @@
     }
 
     /**
+     * Concat an array of {@link NetworkCapabilities.Transport} in string format.
+     *
+     * @param transports an array of connectivity transports
+     * @return a string of the array of transports.
+     */
+    @NonNull
+    public static String connectivityTransportsToString(
+            @NonNull @ConnectivityTransport int[] transports) {
+        return Arrays.stream(transports).mapToObj(DataUtils::connectivityTransportToString)
+                .collect(Collectors.joining("|"));
+    }
+
+    /**
+     * Convert a {@link NetworkCapabilities.Transport} to a string.
+     *
+     * @param transport the connectivity transport
+     * @return the transport in string
+     */
+    @NonNull
+    public static String connectivityTransportToString(
+            @ConnectivityTransport int transport) {
+        return switch (transport) {
+            case NetworkCapabilities.TRANSPORT_CELLULAR -> "CELLULAR";
+            case NetworkCapabilities.TRANSPORT_WIFI -> "WIFI";
+            case NetworkCapabilities.TRANSPORT_BLUETOOTH -> "BLUETOOTH";
+            case NetworkCapabilities.TRANSPORT_ETHERNET -> "ETHERNET";
+            case NetworkCapabilities.TRANSPORT_VPN -> "VPN";
+            case NetworkCapabilities.TRANSPORT_WIFI_AWARE -> "WIFI_AWARE";
+            case NetworkCapabilities.TRANSPORT_LOWPAN -> "LOWPAN";
+            case NetworkCapabilities.TRANSPORT_TEST -> "TEST";
+            case NetworkCapabilities.TRANSPORT_USB -> "USB";
+            case NetworkCapabilities.TRANSPORT_THREAD -> "THREAD";
+            case NetworkCapabilities.TRANSPORT_SATELLITE -> "SATELLITE";
+            default -> "Unknown(" + transport + ")";
+        };
+    }
+
+    /**
      * Convert network capabilities to string.
      *
      * This is for debugging and logging purposes only.
@@ -427,8 +466,7 @@
                                 Arrays.stream(networkRequest.getNativeNetworkRequest()
                                                 .getEnterpriseIds())
                                         .boxed().collect(Collectors.toSet()),
-                                Arrays.stream(networkRequest.getNativeNetworkRequest()
-                                                .getTransportTypes())
+                                Arrays.stream(networkRequest.getTransportTypes())
                                         .boxed().collect(Collectors.toSet())
                                 ),
                         v -> new NetworkRequestList()).add(networkRequest);
diff --git a/src/java/com/android/internal/telephony/data/TelephonyNetworkRequest.java b/src/java/com/android/internal/telephony/data/TelephonyNetworkRequest.java
index b059100..2be3556 100644
--- a/src/java/com/android/internal/telephony/data/TelephonyNetworkRequest.java
+++ b/src/java/com/android/internal/telephony/data/TelephonyNetworkRequest.java
@@ -24,6 +24,7 @@
 import android.net.NetworkRequest;
 import android.net.NetworkSpecifier;
 import android.os.SystemClock;
+import android.telephony.Annotation.ConnectivityTransport;
 import android.telephony.Annotation.NetCapability;
 import android.telephony.data.ApnSetting;
 import android.telephony.data.DataProfile;
@@ -238,6 +239,23 @@
     }
 
     /**
+     * @see NetworkRequest#getTransportTypes()
+     */
+    @NonNull
+    @ConnectivityTransport
+    public int[] getTransportTypes() {
+        return mNativeNetworkRequest.getTransportTypes();
+    }
+
+    /**
+     * @return {@code true} if the request can be served on the specified transport.
+     * @see NetworkRequest#hasTransport
+     */
+    public boolean hasTransport(@ConnectivityTransport int transport) {
+        return mNativeNetworkRequest.hasTransport(transport);
+    }
+
+    /**
      * @see NetworkRequest#canBeSatisfiedBy(NetworkCapabilities)
      */
     public boolean canBeSatisfiedBy(@Nullable NetworkCapabilities nc) {
diff --git a/src/java/com/android/internal/telephony/satellite/DatagramController.java b/src/java/com/android/internal/telephony/satellite/DatagramController.java
index 2f18796..8271187 100644
--- a/src/java/com/android/internal/telephony/satellite/DatagramController.java
+++ b/src/java/com/android/internal/telephony/satellite/DatagramController.java
@@ -105,6 +105,7 @@
     private long mDatagramWaitTimeForConnectedState;
     private long mModemImageSwitchingDuration;
     private boolean mWaitForDeviceAlignmentInDemoDatagram;
+    private long mDatagramWaitTimeForConnectedStateForLastMessage;
     @GuardedBy("mLock")
     @SatelliteManager.SatelliteModemState
     private int mSatelltieModemState = SatelliteManager.SATELLITE_MODEM_STATE_UNKNOWN;
@@ -161,6 +162,8 @@
         mModemImageSwitchingDuration = getSatelliteModemImageSwitchingDurationMillis();
         mWaitForDeviceAlignmentInDemoDatagram =
                 getWaitForDeviceAlignmentInDemoDatagramFromResources();
+        mDatagramWaitTimeForConnectedStateForLastMessage =
+                getDatagramWaitForConnectedStateForLastMessageTimeoutMillis();
         mDemoModeDatagramList = new ArrayList<>();
     }
 
@@ -447,13 +450,17 @@
     }
 
     @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
-    public long getDatagramWaitTimeForConnectedState() {
+    public long getDatagramWaitTimeForConnectedState(boolean isLastSosMessage) {
         synchronized (mLock) {
+            long timeout = isLastSosMessage ? mDatagramWaitTimeForConnectedStateForLastMessage :
+                    mDatagramWaitTimeForConnectedState;
+            logd("getDatagramWaitTimeForConnectedState: isLastSosMessage=" + isLastSosMessage
+                    + ", timeout=" + timeout + ", modemState=" + mSatelltieModemState);
             if (mSatelltieModemState == SATELLITE_MODEM_STATE_OFF
                     || mSatelltieModemState == SATELLITE_MODEM_STATE_IDLE) {
-                return (mDatagramWaitTimeForConnectedState + mModemImageSwitchingDuration);
+                return (timeout + mModemImageSwitchingDuration);
             }
-            return mDatagramWaitTimeForConnectedState;
+            return timeout;
         }
     }
 
@@ -553,6 +560,11 @@
                 R.integer.config_satellite_modem_image_switching_duration_millis);
     }
 
+    private long getDatagramWaitForConnectedStateForLastMessageTimeoutMillis() {
+        return mContext.getResources().getInteger(
+                R.integer.config_datagram_wait_for_connected_state_for_last_message_timeout_millis);
+    }
+
     /**
      * This API can be used by only CTS to override the cached value for the device overlay config
      * value : config_send_satellite_datagram_to_modem_in_demo_mode, which determines whether
diff --git a/src/java/com/android/internal/telephony/satellite/DatagramDispatcher.java b/src/java/com/android/internal/telephony/satellite/DatagramDispatcher.java
index f764b2b..3984a5f 100644
--- a/src/java/com/android/internal/telephony/satellite/DatagramDispatcher.java
+++ b/src/java/com/android/internal/telephony/satellite/DatagramDispatcher.java
@@ -104,6 +104,7 @@
             mPendingNonEmergencyDatagramsMap = new LinkedHashMap<>();
 
     private long mWaitTimeForDatagramSendingResponse;
+    private long mWaitTimeForDatagramSendingForLastMessageResponse;
     @SatelliteManager.DatagramType
     private int mLastSendRequestDatagramType = DATAGRAM_TYPE_UNKNOWN;
 
@@ -142,6 +143,8 @@
             mSendingDatagramInProgress = false;
         }
         mWaitTimeForDatagramSendingResponse = getWaitForDatagramSendingResponseTimeoutMillis();
+        mWaitTimeForDatagramSendingForLastMessageResponse =
+                getWaitForDatagramSendingResponseForLastMessageTimeoutMillis();
     }
 
     private static final class DatagramDispatcherHandlerRequest {
@@ -680,7 +683,8 @@
         }
         sendMessageDelayed(obtainMessage(
                         EVENT_DATAGRAM_WAIT_FOR_CONNECTED_STATE_TIMED_OUT, datagramArgs),
-                mDatagramController.getDatagramWaitTimeForConnectedState());
+                mDatagramController.getDatagramWaitTimeForConnectedState(
+                        SatelliteServiceUtils.isLastSosMessage(datagramArgs.datagramType)));
     }
 
     private void stopDatagramWaitForConnectedStateTimer() {
@@ -709,9 +713,13 @@
             logd("WaitForDatagramSendingResponseTimer was already started");
             return;
         }
+        long waitTime = SatelliteServiceUtils.isLastSosMessage(argument.datagramType)
+                ? mWaitTimeForDatagramSendingForLastMessageResponse
+                : mWaitTimeForDatagramSendingResponse;
+        logd("startWaitForDatagramSendingResponseTimer: datagramType=" + argument.datagramType
+                + ", waitTime=" + waitTime);
         sendMessageDelayed(obtainMessage(
-                EVENT_WAIT_FOR_DATAGRAM_SENDING_RESPONSE_TIMED_OUT, argument),
-                mWaitTimeForDatagramSendingResponse);
+                EVENT_WAIT_FOR_DATAGRAM_SENDING_RESPONSE_TIMED_OUT, argument), waitTime);
     }
 
     private void stopWaitForDatagramSendingResponseTimer() {
@@ -760,6 +768,11 @@
                 R.integer.config_wait_for_datagram_sending_response_timeout_millis);
     }
 
+    private long getWaitForDatagramSendingResponseForLastMessageTimeoutMillis() {
+        return mContext.getResources().getInteger(R.integer
+                .config_wait_for_datagram_sending_response_for_last_message_timeout_millis);
+    }
+
     private boolean shouldProcessEventSendSatelliteDatagramDone(
             @NonNull SendSatelliteDatagramArgument argument) {
         synchronized (mLock) {
diff --git a/src/java/com/android/internal/telephony/satellite/DatagramReceiver.java b/src/java/com/android/internal/telephony/satellite/DatagramReceiver.java
index f1f0fde..21e4318 100644
--- a/src/java/com/android/internal/telephony/satellite/DatagramReceiver.java
+++ b/src/java/com/android/internal/telephony/satellite/DatagramReceiver.java
@@ -818,7 +818,7 @@
         }
         sendMessageDelayed(obtainMessage(
                         EVENT_DATAGRAM_WAIT_FOR_CONNECTED_STATE_TIMED_OUT),
-                mDatagramController.getDatagramWaitTimeForConnectedState());
+                mDatagramController.getDatagramWaitTimeForConnectedState(false));
     }
 
     private void stopDatagramWaitForConnectedStateTimer() {
diff --git a/src/java/com/android/internal/telephony/satellite/SatelliteController.java b/src/java/com/android/internal/telephony/satellite/SatelliteController.java
index bde255b..7723bbd 100644
--- a/src/java/com/android/internal/telephony/satellite/SatelliteController.java
+++ b/src/java/com/android/internal/telephony/satellite/SatelliteController.java
@@ -155,6 +155,7 @@
     private static final boolean DEBUG = !"user".equals(Build.TYPE);
     /** File used to store shared preferences related to satellite. */
     public static final String SATELLITE_SHARED_PREF = "satellite_shared_pref";
+    public static final String SATELLITE_SUBSCRIPTION_ID = "satellite_subscription_id";
     /** Value to pass for the setting key SATELLITE_MODE_ENABLED, enabled = 1, disabled = 0 */
     public static final int SATELLITE_MODE_ENABLED_TRUE = 1;
     public static final int SATELLITE_MODE_ENABLED_FALSE = 0;
@@ -1645,6 +1646,7 @@
                 /* We have already successfully queried the satellite modem. */
                 Bundle bundle = new Bundle();
                 bundle.putBoolean(SatelliteManager.KEY_SATELLITE_SUPPORTED, mIsSatelliteSupported);
+                bundle.putInt(SATELLITE_SUBSCRIPTION_ID, subId);
                 result.send(SATELLITE_RESULT_SUCCESS, bundle);
                 return;
             }
diff --git a/src/java/com/android/internal/telephony/satellite/SatelliteModemInterface.java b/src/java/com/android/internal/telephony/satellite/SatelliteModemInterface.java
index 58260d1..8a26fd2 100644
--- a/src/java/com/android/internal/telephony/satellite/SatelliteModemInterface.java
+++ b/src/java/com/android/internal/telephony/satellite/SatelliteModemInterface.java
@@ -66,6 +66,8 @@
     @NonNull private static SatelliteModemInterface sInstance;
     @NonNull private final Context mContext;
     @NonNull private final DemoSimulator mDemoSimulator;
+    @NonNull private final SatelliteListener mVendorListener;
+    @NonNull private final SatelliteListener mDemoListener;
     @VisibleForTesting(visibility = VisibleForTesting.Visibility.PRIVATE)
     @NonNull protected final ExponentialBackoff mExponentialBackoff;
     @NonNull private final Object mLock = new Object();
@@ -237,6 +239,8 @@
             SatelliteController satelliteController, @NonNull Looper looper) {
         mContext = context;
         mDemoSimulator = DemoSimulator.make(context, satelliteController);
+        mVendorListener = new SatelliteListener(false);
+        mDemoListener = new SatelliteListener(true);
         mIsSatelliteServiceSupported = getSatelliteServiceSupport();
         mSatelliteController = satelliteController;
         mExponentialBackoff = new ExponentialBackoff(REBIND_INITIAL_DELAY, REBIND_MAXIMUM_DELAY,
@@ -349,11 +353,8 @@
             mSatelliteService = ISatellite.Stub.asInterface(service);
             mExponentialBackoff.stop();
             try {
-                SatelliteListener vendorListener = new SatelliteListener(false);
-                mSatelliteService.setSatelliteListener(vendorListener);
-
-                SatelliteListener demoListener = new SatelliteListener(true);
-                mDemoSimulator.setSatelliteListener(demoListener);
+                mSatelliteService.setSatelliteListener(mVendorListener);
+                mDemoSimulator.setSatelliteListener(mDemoListener);
             } catch (RemoteException e) {
                 // TODO: Retry setSatelliteListener
                 logd("setSatelliteListener: RemoteException " + e);
diff --git a/src/java/com/android/internal/telephony/satellite/SatelliteServiceUtils.java b/src/java/com/android/internal/telephony/satellite/SatelliteServiceUtils.java
index 67b2586..d33fd73 100644
--- a/src/java/com/android/internal/telephony/satellite/SatelliteServiceUtils.java
+++ b/src/java/com/android/internal/telephony/satellite/SatelliteServiceUtils.java
@@ -347,7 +347,9 @@
     }
 
     /**
-     *
+     * Check if the datagramType is the sos message (DATAGRAM_TYPE_SOS_MESSAGE,
+     * DATAGRAM_TYPE_LAST_SOS_MESSAGE_STILL_NEED_HELP,
+     * DATAGRAM_TYPE_LAST_SOS_MESSAGE_NO_HELP_NEEDED) or not
      */
     public static boolean isSosMessage(int datagramType) {
         return datagramType == SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE
@@ -356,6 +358,16 @@
     }
 
     /**
+     * Check if the datagramType is the last sos message
+     * (DATAGRAM_TYPE_LAST_SOS_MESSAGE_STILL_NEED_HELP or
+     * DATAGRAM_TYPE_LAST_SOS_MESSAGE_NO_HELP_NEEDED) or not
+     */
+    public static boolean isLastSosMessage(int datagramType) {
+        return datagramType == SatelliteManager.DATAGRAM_TYPE_LAST_SOS_MESSAGE_STILL_NEED_HELP
+                || datagramType == SatelliteManager.DATAGRAM_TYPE_LAST_SOS_MESSAGE_NO_HELP_NEEDED;
+    }
+
+    /**
      * Return phone associated with phoneId 0.
      *
      * @return phone associated with phoneId 0 or {@code null} if it doesn't exist.
diff --git a/tests/telephonytests/src/com/android/internal/telephony/data/DataNetworkControllerTest.java b/tests/telephonytests/src/com/android/internal/telephony/data/DataNetworkControllerTest.java
index d11b730..7f42ad0 100644
--- a/tests/telephonytests/src/com/android/internal/telephony/data/DataNetworkControllerTest.java
+++ b/tests/telephonytests/src/com/android/internal/telephony/data/DataNetworkControllerTest.java
@@ -138,6 +138,8 @@
 import java.util.Set;
 import java.util.concurrent.Executor;
 
+import javax.annotation.Nullable;
+
 @RunWith(AndroidTestingRunner.class)
 @TestableLooper.RunWithLooper
 public class DataNetworkControllerTest extends TelephonyTest {
@@ -1138,18 +1140,23 @@
     }
 
     private @NonNull TelephonyNetworkRequest createNetworkRequest(Integer... capabilities) {
-        return createNetworkRequest(false, capabilities);
+        return createNetworkRequest(null, capabilities);
     }
 
-    private @NonNull TelephonyNetworkRequest createNetworkRequest(boolean restricted,
+    private @NonNull TelephonyNetworkRequest createNetworkRequest(@Nullable Boolean restricted,
                                                                   Integer... capabilities) {
         NetworkCapabilities netCaps = new NetworkCapabilities();
         for (int networkCapability : capabilities) {
             netCaps.addCapability(networkCapability);
         }
 
-        if (restricted) {
-            netCaps.removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED);
+        if (restricted != null) {
+            if (restricted) {
+                netCaps.removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED);
+            }
+        } else {
+            // Data Network uses the same to define its own capabilities.
+            netCaps.maybeMarkCapabilitiesRestricted();
         }
 
         NetworkRequest nativeNetworkRequest = new NetworkRequest(netCaps,
@@ -1741,6 +1748,14 @@
         verify(mMockedDataNetworkControllerCallback, never())
                 .onConnectedInternetDataNetworksChanged(any());
 
+        // However, WLAN network setup shouldn't be affected
+        doReturn(AccessNetworkConstants.TRANSPORT_TYPE_WLAN).when(mAccessNetworksManager)
+                .getPreferredTransportByNetworkCapability(anyInt());
+        mDataNetworkControllerUT.obtainMessage(5 /*EVENT_REEVALUATE_UNSATISFIED_NETWORK_REQUESTS*/,
+                DataEvaluation.DataEvaluationReason.PREFERRED_TRANSPORT_CHANGED).sendToTarget();
+        processAllMessages();
+        verify(mMockedDataNetworkControllerCallback).onConnectedInternetDataNetworksChanged(any());
+
         mIsNonTerrestrialNetwork = false;
     }
 
@@ -3319,12 +3334,20 @@
                 NetworkCapabilities.NET_CAPABILITY_IMS);
         TelephonyNetworkRequest secondTnr = createNetworkRequest(
                 NetworkCapabilities.NET_CAPABILITY_IMS);
+        TelephonyNetworkRequest thirdTnr = new TelephonyNetworkRequest(
+                new NetworkRequest((new NetworkCapabilities.Builder())
+                        .addCapability(NetworkCapabilities.NET_CAPABILITY_IMS)
+                        .addTransportType(NetworkCapabilities.TRANSPORT_SATELLITE)
+                        .build(),
+                        ConnectivityManager.TYPE_MOBILE, ++mNetworkRequestId,
+                        NetworkRequest.Type.REQUEST), mPhone, mFeatureFlags);
 
         mDataNetworkControllerUT.addNetworkRequest(firstTnr);
         processAllMessages();
 
         mDataNetworkControllerUT.removeNetworkRequest(firstTnr);
         mDataNetworkControllerUT.addNetworkRequest(secondTnr);
+        mDataNetworkControllerUT.addNetworkRequest(thirdTnr);
         processAllFutureMessages();
 
         verify(mMockedWwanDataServiceManager, times(1)).setupDataCall(anyInt(),
@@ -4157,6 +4180,7 @@
 
         NetworkCapabilities netCaps = new NetworkCapabilities();
         netCaps.addCapability(NetworkCapabilities.NET_CAPABILITY_IMS);
+        netCaps.maybeMarkCapabilitiesRestricted();
         netCaps.setRequestorPackageName(FAKE_MMTEL_PACKAGE);
 
         NetworkRequest nativeNetworkRequest = new NetworkRequest(netCaps,
@@ -4209,6 +4233,7 @@
         // setup emergency data network.
         NetworkCapabilities netCaps = new NetworkCapabilities();
         netCaps.addCapability(NetworkCapabilities.NET_CAPABILITY_EIMS);
+        netCaps.maybeMarkCapabilitiesRestricted();
         netCaps.setRequestorPackageName(FAKE_MMTEL_PACKAGE);
 
         NetworkRequest nativeNetworkRequest = new NetworkRequest(netCaps,
@@ -5097,7 +5122,8 @@
         serviceStateChanged(TelephonyManager.NETWORK_TYPE_LTE,
                 NetworkRegistrationInfo.REGISTRATION_STATE_HOME);
         mDataNetworkControllerUT.addNetworkRequest(
-                createNetworkRequest(NetworkCapabilities.NET_CAPABILITY_INTERNET));
+                createNetworkRequest(true/*restricted*/,
+                        NetworkCapabilities.NET_CAPABILITY_INTERNET));
         processAllMessages();
         verifyConnectedNetworkHasDataProfile(mEsimBootstrapDataProfile);
 
@@ -5108,13 +5134,6 @@
                 .get(AccessNetworkConstants.TRANSPORT_TYPE_WWAN), 2);
         processAllMessages();
         verifyConnectedNetworkHasDataProfile(mEsimBootstrapImsProfile);
-
-        serviceStateChanged(TelephonyManager.NETWORK_TYPE_LTE,
-                NetworkRegistrationInfo.REGISTRATION_STATE_HOME);
-        mDataNetworkControllerUT.addNetworkRequest(
-                createNetworkRequest(NetworkCapabilities.NET_CAPABILITY_RCS));
-        processAllMessages();
-        verifyNoConnectedNetworkHasCapability(NetworkCapabilities.NET_CAPABILITY_RCS);
     }
 
     @Test
@@ -5153,7 +5172,7 @@
         serviceStateChanged(TelephonyManager.NETWORK_TYPE_LTE,
                 NetworkRegistrationInfo.REGISTRATION_STATE_HOME);
         mDataNetworkControllerUT.addNetworkRequest(
-                createNetworkRequest(NetworkCapabilities.NET_CAPABILITY_INTERNET));
+                createNetworkRequest(true, NetworkCapabilities.NET_CAPABILITY_INTERNET));
         processAllMessages();
         // With allowed data limit unlimited, connection is allowed
         verifyConnectedNetworkHasDataProfile(mEsimBootstrapDataProfile);
diff --git a/tests/telephonytests/src/com/android/internal/telephony/satellite/DatagramDispatcherTest.java b/tests/telephonytests/src/com/android/internal/telephony/satellite/DatagramDispatcherTest.java
index 4896671..dba288e 100644
--- a/tests/telephonytests/src/com/android/internal/telephony/satellite/DatagramDispatcherTest.java
+++ b/tests/telephonytests/src/com/android/internal/telephony/satellite/DatagramDispatcherTest.java
@@ -98,6 +98,12 @@
     private static final long TEST_DATAGRAM_WAIT_FOR_CONNECTED_STATE_TIMEOUT_MILLIS =
             TimeUnit.SECONDS.toMillis(60);
     private static final Long TIMEOUT_DATAGRAM_DELAY_IN_DEMO_MODE = TimeUnit.SECONDS.toMillis(10);
+    private static final long
+            TEST_DATAGRAM_WAIT_FOR_CONNECTED_STATE_FOR_LAST_MESSAGE_TIMEOUT_MILLIS =
+            TimeUnit.SECONDS.toMillis(60);
+    private static final int
+            TEST_WAIT_FOR_DATAGRAM_SENDING_RESPONSE_FOR_LAST_MESSAGE_TIMEOUT_MILLIS =
+            (int) TimeUnit.SECONDS.toMillis(60);
 
     private TestDatagramDispatcher mDatagramDispatcherUT;
 
@@ -188,8 +194,11 @@
             clearInvocations(mMockSatelliteModemInterface);
             doReturn(true).when(mMockDatagramController)
                     .needsWaitingForSatelliteConnected(eq(datagramType));
-            when(mMockDatagramController.getDatagramWaitTimeForConnectedState())
+            when(mMockDatagramController.getDatagramWaitTimeForConnectedState(eq(false)))
                     .thenReturn(TEST_DATAGRAM_WAIT_FOR_CONNECTED_STATE_TIMEOUT_MILLIS);
+            when(mMockDatagramController.getDatagramWaitTimeForConnectedState(eq(true)))
+                    .thenReturn(
+                            TEST_DATAGRAM_WAIT_FOR_CONNECTED_STATE_FOR_LAST_MESSAGE_TIMEOUT_MILLIS);
             mResultListener.clear();
 
             mDatagramDispatcherUT.sendSatelliteDatagram(SUB_ID, datagramType, mDatagram,
@@ -201,7 +210,8 @@
                     eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT),
                     eq(1),
                     eq(SATELLITE_RESULT_SUCCESS));
-            mInOrder.verify(mMockDatagramController).getDatagramWaitTimeForConnectedState();
+            mInOrder.verify(mMockDatagramController).getDatagramWaitTimeForConnectedState(
+                    eq(SatelliteServiceUtils.isLastSosMessage(datagramType)));
             verifyZeroInteractions(mMockSatelliteModemInterface);
             assertTrue(mDatagramDispatcherUT.isDatagramWaitForConnectedStateTimerStarted());
 
@@ -247,7 +257,8 @@
             verifyZeroInteractions(mMockSatelliteModemInterface);
             mInOrder.verify(mMockDatagramController)
                     .needsWaitingForSatelliteConnected(eq(datagramType));
-            mInOrder.verify(mMockDatagramController).getDatagramWaitTimeForConnectedState();
+            mInOrder.verify(mMockDatagramController).getDatagramWaitTimeForConnectedState(
+                    eq(SatelliteServiceUtils.isLastSosMessage(datagramType)));
             assertTrue(mDatagramDispatcherUT.isDatagramWaitForConnectedStateTimerStarted());
 
             moveTimeForward(TEST_DATAGRAM_WAIT_FOR_CONNECTED_STATE_TIMEOUT_MILLIS);
@@ -287,7 +298,8 @@
             verifyZeroInteractions(mMockSatelliteModemInterface);
             mInOrder.verify(mMockDatagramController)
                     .needsWaitingForSatelliteConnected(eq(datagramType));
-            mInOrder.verify(mMockDatagramController).getDatagramWaitTimeForConnectedState();
+            mInOrder.verify(mMockDatagramController).getDatagramWaitTimeForConnectedState(
+                    eq(SatelliteServiceUtils.isLastSosMessage(datagramType)));
             assertTrue(mDatagramDispatcherUT.isDatagramWaitForConnectedStateTimerStarted());
             assertEquals(0, mResultListener.size());
 
@@ -305,86 +317,102 @@
 
     @Test
     public void testSendSatelliteDatagram_timeout() throws  Exception {
-        doAnswer(invocation -> {
-            Message message = (Message) invocation.getArguments()[3];
-
-            mDatagramDispatcherUT.obtainMessage(2 /*EVENT_SEND_SATELLITE_DATAGRAM_DONE*/,
-                            new AsyncResult(message.obj, null, null))
-                    .sendToTarget();
-
-            // DatagramDispatcher should ignore the second EVENT_SEND_SATELLITE_DATAGRAM_DONE
-            mDatagramDispatcherUT.obtainMessage(2 /*EVENT_SEND_SATELLITE_DATAGRAM_DONE*/,
-                            new AsyncResult(message.obj, null, null))
-                    .sendToTarget();
-
-            return null;
-        }).when(mMockSatelliteModemInterface).sendSatelliteDatagram(any(SatelliteDatagram.class),
-                anyBoolean(), anyBoolean(), any(Message.class));
-        doReturn(false).when(mMockDatagramController)
-                .needsWaitingForSatelliteConnected(eq(DATAGRAM_TYPE1));
-        when(mMockDatagramController.getDatagramWaitTimeForConnectedState())
+        when(mMockDatagramController.getDatagramWaitTimeForConnectedState(eq(false)))
                 .thenReturn(TEST_DATAGRAM_WAIT_FOR_CONNECTED_STATE_TIMEOUT_MILLIS);
+        when(mMockDatagramController.getDatagramWaitTimeForConnectedState(eq(true)))
+                .thenReturn(TEST_DATAGRAM_WAIT_FOR_CONNECTED_STATE_FOR_LAST_MESSAGE_TIMEOUT_MILLIS);
         mContextFixture.putIntResource(
                 R.integer.config_wait_for_datagram_sending_response_timeout_millis,
                 TEST_WAIT_FOR_DATAGRAM_SENDING_RESPONSE_TIMEOUT_MILLIS);
+        mContextFixture.putIntResource(
+                R.integer.config_wait_for_datagram_sending_response_for_last_message_timeout_millis,
+                TEST_WAIT_FOR_DATAGRAM_SENDING_RESPONSE_FOR_LAST_MESSAGE_TIMEOUT_MILLIS);
         mResultListener.clear();
+        int[] sosDatagramTypes = {DATAGRAM_TYPE1, DATAGRAM_TYPE4, DATAGRAM_TYPE5};
+        for (int datagramType : sosDatagramTypes) {
+            doAnswer(invocation -> {
+                Message message = (Message) invocation.getArguments()[3];
 
-        mDatagramDispatcherUT.sendSatelliteDatagram(SUB_ID, DATAGRAM_TYPE1, mDatagram,
-                true, mResultListener::offer);
-        processAllMessages();
-        mInOrder.verify(mMockDatagramController)
-                .needsWaitingForSatelliteConnected(eq(DATAGRAM_TYPE1));
-        mInOrder.verify(mMockDatagramController).isPollingInIdleState();
-        mInOrder.verify(mMockDatagramController)
-                .updateSendStatus(eq(SUB_ID), eq(DATAGRAM_TYPE1),
-                        eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING), eq(1),
-                        eq(SATELLITE_RESULT_SUCCESS));
-        mInOrder.verify(mMockDatagramController)
-                .updateSendStatus(eq(SUB_ID), eq(DATAGRAM_TYPE1),
-                        eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_SEND_SUCCESS), eq(0),
-                        eq(SATELLITE_RESULT_SUCCESS));
-        mInOrder.verify(mMockDatagramController)
-                .updateSendStatus(eq(SUB_ID), eq(DATAGRAM_TYPE1),
-                        eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE), eq(0),
-                        eq(SATELLITE_RESULT_SUCCESS));
-        verifyNoMoreInteractions(mMockDatagramController);
-        verify(mMockSatelliteModemInterface, times(1)).sendSatelliteDatagram(
-                any(SatelliteDatagram.class), anyBoolean(), anyBoolean(), any(Message.class));
-        assertThat(mResultListener.peek()).isEqualTo(SATELLITE_RESULT_SUCCESS);
-        verify(mMockSessionMetricsStats, times(1)).addCountOfSuccessfulOutgoingDatagram();
-        clearInvocations(mMockSatelliteModemInterface);
-        clearInvocations(mMockDatagramController);
-        clearInvocations(mMockSessionMetricsStats);
-        mResultListener.clear();
+                mDatagramDispatcherUT.obtainMessage(2 /*EVENT_SEND_SATELLITE_DATAGRAM_DONE*/,
+                                new AsyncResult(message.obj, null, null))
+                        .sendToTarget();
 
-        // No response for the send request from modem
-        doNothing().when(mMockSatelliteModemInterface).sendSatelliteDatagram(
-                any(SatelliteDatagram.class), anyBoolean(), anyBoolean(), any(Message.class));
+                // DatagramDispatcher should ignore the second EVENT_SEND_SATELLITE_DATAGRAM_DONE
+                mDatagramDispatcherUT.obtainMessage(2 /*EVENT_SEND_SATELLITE_DATAGRAM_DONE*/,
+                                new AsyncResult(message.obj, null, null))
+                        .sendToTarget();
 
-        mDatagramDispatcherUT.sendSatelliteDatagram(SUB_ID, DATAGRAM_TYPE1, mDatagram,
-                true, mResultListener::offer);
-        processAllMessages();
-        mInOrder.verify(mMockDatagramController)
-                .needsWaitingForSatelliteConnected(eq(DATAGRAM_TYPE1));
-        mInOrder.verify(mMockDatagramController).isPollingInIdleState();
-        mInOrder.verify(mMockDatagramController)
-                .updateSendStatus(eq(SUB_ID), eq(DATAGRAM_TYPE1),
-                        eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING), eq(1),
-                        eq(SATELLITE_RESULT_SUCCESS));
-        mInOrder.verify(mMockDatagramController)
-                .updateSendStatus(eq(SUB_ID), eq(DATAGRAM_TYPE1),
-                        eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_SEND_FAILED), eq(1),
-                        eq(SATELLITE_RESULT_MODEM_TIMEOUT));
-        mInOrder.verify(mMockDatagramController)
-                .updateSendStatus(eq(SUB_ID), eq(DATAGRAM_TYPE1),
-                        eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE), eq(0),
-                        eq(SATELLITE_RESULT_SUCCESS));
-        verifyNoMoreInteractions(mMockDatagramController);
-        verify(mMockSatelliteModemInterface, times(1)).sendSatelliteDatagram(
-                any(SatelliteDatagram.class), anyBoolean(), anyBoolean(), any(Message.class));
-        verify(mMockSatelliteModemInterface).abortSendingSatelliteDatagrams(any(Message.class));
-        assertThat(mResultListener.peek()).isEqualTo(SATELLITE_RESULT_MODEM_TIMEOUT);
-        verify(mMockSessionMetricsStats, times(1)).addCountOfFailedOutgoingDatagram();
+                return null;
+            }).when(mMockSatelliteModemInterface).sendSatelliteDatagram(
+                    any(SatelliteDatagram.class),
+                    anyBoolean(), anyBoolean(), any(Message.class));
+            doReturn(false).when(mMockDatagramController)
+                    .needsWaitingForSatelliteConnected(eq(datagramType));
+
+            mDatagramDispatcherUT.sendSatelliteDatagram(SUB_ID, datagramType, mDatagram,
+                    true, mResultListener::offer);
+            processAllMessages();
+            mInOrder.verify(mMockDatagramController)
+                    .needsWaitingForSatelliteConnected(eq(datagramType));
+            mInOrder.verify(mMockDatagramController).isPollingInIdleState();
+            mInOrder.verify(mMockDatagramController)
+                    .updateSendStatus(eq(SUB_ID), eq(datagramType),
+                            eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING), eq(1),
+                            eq(SATELLITE_RESULT_SUCCESS));
+            mInOrder.verify(mMockDatagramController)
+                    .updateSendStatus(eq(SUB_ID), eq(datagramType),
+                            eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_SEND_SUCCESS),
+                            eq(0),
+                            eq(SATELLITE_RESULT_SUCCESS));
+            mInOrder.verify(mMockDatagramController)
+                    .updateSendStatus(eq(SUB_ID), eq(datagramType),
+                            eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE), eq(0),
+                            eq(SATELLITE_RESULT_SUCCESS));
+            verifyNoMoreInteractions(mMockDatagramController);
+            verify(mMockSatelliteModemInterface, times(1)).sendSatelliteDatagram(
+                    any(SatelliteDatagram.class), anyBoolean(), anyBoolean(), any(Message.class));
+            assertThat(mResultListener.peek()).isEqualTo(SATELLITE_RESULT_SUCCESS);
+            verify(mMockSessionMetricsStats, times(1)).addCountOfSuccessfulOutgoingDatagram();
+            clearInvocations(mMockSatelliteModemInterface);
+            clearInvocations(mMockDatagramController);
+            clearInvocations(mMockSessionMetricsStats);
+            mResultListener.clear();
+
+            // No response for the send request from modem
+            doNothing().when(mMockSatelliteModemInterface).sendSatelliteDatagram(
+                    any(SatelliteDatagram.class), anyBoolean(), anyBoolean(), any(Message.class));
+
+            mDatagramDispatcherUT.sendSatelliteDatagram(SUB_ID, datagramType, mDatagram,
+                    true, mResultListener::offer);
+            processAllMessages();
+            mInOrder.verify(mMockDatagramController)
+                    .needsWaitingForSatelliteConnected(eq(datagramType));
+            mInOrder.verify(mMockDatagramController).isPollingInIdleState();
+            mInOrder.verify(mMockDatagramController)
+                    .updateSendStatus(eq(SUB_ID), eq(datagramType),
+                            eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING), eq(1),
+                            eq(SATELLITE_RESULT_SUCCESS));
+            mInOrder.verify(mMockDatagramController)
+                    .updateSendStatus(eq(SUB_ID), eq(datagramType),
+                            eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_SEND_FAILED),
+                            eq(1),
+                            eq(SATELLITE_RESULT_MODEM_TIMEOUT));
+            mInOrder.verify(mMockDatagramController)
+                    .updateSendStatus(eq(SUB_ID), eq(datagramType),
+                            eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE), eq(0),
+                            eq(SATELLITE_RESULT_SUCCESS));
+            verifyNoMoreInteractions(mMockDatagramController);
+            verify(mMockSatelliteModemInterface, times(1)).sendSatelliteDatagram(
+                    any(SatelliteDatagram.class), anyBoolean(), anyBoolean(), any(Message.class));
+            verify(mMockSatelliteModemInterface).abortSendingSatelliteDatagrams(any(Message.class));
+            assertThat(mResultListener.peek()).isEqualTo(SATELLITE_RESULT_MODEM_TIMEOUT);
+            verify(mMockSessionMetricsStats, times(1)).addCountOfFailedOutgoingDatagram();
+
+            clearInvocations(mMockSatelliteModemInterface);
+            clearInvocations(mMockDatagramController);
+            clearInvocations(mMockSessionMetricsStats);
+            mResultListener.clear();
+        }
     }
 
     @Test
@@ -629,87 +657,96 @@
     }
 
     @Test
-    public void testSendSatelliteDatagramToModemInDemoMode()
-            throws Exception {
+    public void testSendSatelliteDatagramToModemInDemoMode() throws Exception {
         when(mFeatureFlags.oemEnabledSatelliteFlag()).thenReturn(true);
 
-        doAnswer(invocation -> {
-            Message message = (Message) invocation.getArguments()[3];
-            mDatagramDispatcherUT.obtainMessage(2 /*EVENT_SEND_SATELLITE_DATAGRAM_DONE*/,
-                            new AsyncResult(message.obj, null, null))
-                    .sendToTarget();
-            return null;
-        }).when(mMockSatelliteModemInterface).sendSatelliteDatagram(any(SatelliteDatagram.class),
-                anyBoolean(), anyBoolean(), any(Message.class));
         mDatagramDispatcherUT.setDemoMode(true);
         mDatagramDispatcherUT.setDeviceAlignedWithSatellite(true);
-        mIntegerConsumerSemaphore.drainPermits();
 
-        // Test when overlay config config_send_satellite_datagram_to_modem_in_demo_mode is true
-        mDatagramDispatcherUT.setShouldSendDatagramToModemInDemoMode(null);
-        mContextFixture.putBooleanResource(mConfigSendSatelliteDatagramToModemInDemoMode, true);
-        mDatagramDispatcherUT.sendSatelliteDatagram(SUB_ID, DATAGRAM_TYPE1, mDatagram,
-                true, mIntegerConsumer);
-        processAllMessages();
-        moveTimeForward(TIMEOUT_DATAGRAM_DELAY_IN_DEMO_MODE);
-        processAllMessages();
-        waitForIntegerConsumerResult(1);
-        assertEquals(SATELLITE_RESULT_SUCCESS,
-                (int) mIntegerConsumerResult.get(0));
-        mIntegerConsumerResult.clear();
-        verify(mMockSatelliteModemInterface, times(1)).sendSatelliteDatagram(
-                any(SatelliteDatagram.class), anyBoolean(), anyBoolean(), any(Message.class));
+        int[] sosDatagramTypes = {DATAGRAM_TYPE1, DATAGRAM_TYPE4, DATAGRAM_TYPE5};
+        for (int datagramType : sosDatagramTypes) {
+            mIntegerConsumerSemaphore.drainPermits();
+            mIntegerConsumerResult.clear();
+            clearInvocations(mMockDatagramController);
+            clearInvocations(mMockSatelliteModemInterface);
+            clearInvocations(mMockSessionMetricsStats);
+            doAnswer(invocation -> {
+                Message message = (Message) invocation.getArguments()[3];
+                mDatagramDispatcherUT.obtainMessage(2 /*EVENT_SEND_SATELLITE_DATAGRAM_DONE*/,
+                                new AsyncResult(message.obj, null, null))
+                        .sendToTarget();
+                return null;
+            }).when(mMockSatelliteModemInterface).sendSatelliteDatagram(
+                    any(SatelliteDatagram.class),
+                    anyBoolean(), anyBoolean(), any(Message.class));
 
-        moveTimeForward(TIMEOUT_DATAGRAM_DELAY_IN_DEMO_MODE);
-        processAllMessages();
-        verify(mMockDatagramController).pushDemoModeDatagram(
-                anyInt(), any(SatelliteDatagram.class));
-        verify(mMockDatagramController).pollPendingSatelliteDatagrams(anyInt(), any());
-        verify(mMockSessionMetricsStats, times(1)).addCountOfSuccessfulOutgoingDatagram();
+            // Test when overlay config config_send_satellite_datagram_to_modem_in_demo_mode is true
+            mDatagramDispatcherUT.setShouldSendDatagramToModemInDemoMode(null);
+            mContextFixture.putBooleanResource(mConfigSendSatelliteDatagramToModemInDemoMode, true);
+            mDatagramDispatcherUT.sendSatelliteDatagram(SUB_ID, datagramType, mDatagram,
+                    true, mIntegerConsumer);
+            processAllMessages();
+            moveTimeForward(TIMEOUT_DATAGRAM_DELAY_IN_DEMO_MODE);
+            processAllMessages();
+            waitForIntegerConsumerResult(1);
+            assertEquals(SATELLITE_RESULT_SUCCESS, (int) mIntegerConsumerResult.get(0));
+            mIntegerConsumerResult.clear();
+            verify(mMockSatelliteModemInterface, times(1)).sendSatelliteDatagram(
+                    any(SatelliteDatagram.class), anyBoolean(), anyBoolean(), any(Message.class));
 
-        // Test when overlay config config_send_satellite_datagram_to_modem_in_demo_mode is false
-        reset(mMockSatelliteModemInterface);
-        mDatagramDispatcherUT.setShouldSendDatagramToModemInDemoMode(null);
-        mContextFixture.putBooleanResource(mConfigSendSatelliteDatagramToModemInDemoMode, false);
-        mDatagramDispatcherUT.sendSatelliteDatagram(SUB_ID, DATAGRAM_TYPE1, mDatagram,
-                true, mIntegerConsumer);
-        processAllMessages();
-        moveTimeForward(TIMEOUT_DATAGRAM_DELAY_IN_DEMO_MODE);
-        processAllMessages();
+            moveTimeForward(TIMEOUT_DATAGRAM_DELAY_IN_DEMO_MODE);
+            processAllMessages();
+            verify(mMockDatagramController).pushDemoModeDatagram(
+                    anyInt(), any(SatelliteDatagram.class));
+            verify(mMockDatagramController).pollPendingSatelliteDatagrams(anyInt(), any());
+            verify(mMockSessionMetricsStats, times(1)).addCountOfSuccessfulOutgoingDatagram();
 
-        waitForIntegerConsumerResult(1);
-        assertEquals(SATELLITE_RESULT_SUCCESS,
-                (int) mIntegerConsumerResult.get(0));
-        mIntegerConsumerResult.clear();
-        verify(mMockSatelliteModemInterface, never()).sendSatelliteDatagram(
-                any(SatelliteDatagram.class), anyBoolean(), anyBoolean(), any(Message.class));
+            // Test when overlay config config_send_satellite_datagram_to_modem_in_demo_mode is
+            // false
+            reset(mMockSatelliteModemInterface);
+            mDatagramDispatcherUT.setShouldSendDatagramToModemInDemoMode(null);
+            mContextFixture.putBooleanResource(mConfigSendSatelliteDatagramToModemInDemoMode,
+                    false);
+            mDatagramDispatcherUT.sendSatelliteDatagram(SUB_ID, datagramType, mDatagram,
+                    true, mIntegerConsumer);
+            processAllMessages();
+            moveTimeForward(TIMEOUT_DATAGRAM_DELAY_IN_DEMO_MODE);
+            processAllMessages();
 
-        moveTimeForward(TIMEOUT_DATAGRAM_DELAY_IN_DEMO_MODE);
-        processAllMessages();
-        verify(mMockDatagramController, times(2)).pushDemoModeDatagram(
-                anyInt(), any(SatelliteDatagram.class));
-        verify(mMockDatagramController, times(2)).pollPendingSatelliteDatagrams(anyInt(), any());
+            waitForIntegerConsumerResult(1);
+            assertEquals(SATELLITE_RESULT_SUCCESS, (int) mIntegerConsumerResult.get(0));
+            mIntegerConsumerResult.clear();
+            verify(mMockSatelliteModemInterface, never()).sendSatelliteDatagram(
+                    any(SatelliteDatagram.class), anyBoolean(), anyBoolean(), any(Message.class));
 
-        // Send datagram one more time
-        reset(mMockSatelliteModemInterface);
-        mDatagramDispatcherUT.sendSatelliteDatagram(SUB_ID, DATAGRAM_TYPE1, mDatagram,
-                true, mIntegerConsumer);
-        processAllMessages();
-        moveTimeForward(TIMEOUT_DATAGRAM_DELAY_IN_DEMO_MODE);
-        processAllMessages();
+            moveTimeForward(TIMEOUT_DATAGRAM_DELAY_IN_DEMO_MODE);
+            processAllMessages();
+            verify(mMockDatagramController, times(2)).pushDemoModeDatagram(
+                    anyInt(), any(SatelliteDatagram.class));
+            verify(mMockDatagramController, times(2)).pollPendingSatelliteDatagrams(anyInt(),
+                    any());
 
-        waitForIntegerConsumerResult(1);
-        assertEquals(SATELLITE_RESULT_SUCCESS,
-                (int) mIntegerConsumerResult.get(0));
-        mIntegerConsumerResult.clear();
-        verify(mMockSatelliteModemInterface, never()).sendSatelliteDatagram(
-                any(SatelliteDatagram.class), anyBoolean(), anyBoolean(), any(Message.class));
+            // Send datagram one more time
+            reset(mMockSatelliteModemInterface);
+            mDatagramDispatcherUT.sendSatelliteDatagram(SUB_ID, datagramType, mDatagram,
+                    true, mIntegerConsumer);
+            processAllMessages();
+            moveTimeForward(TIMEOUT_DATAGRAM_DELAY_IN_DEMO_MODE);
+            processAllMessages();
 
-        moveTimeForward(TIMEOUT_DATAGRAM_DELAY_IN_DEMO_MODE);
-        processAllMessages();
-        verify(mMockDatagramController, times(3)).pushDemoModeDatagram(
-                anyInt(), any(SatelliteDatagram.class));
-        verify(mMockDatagramController, times(3)).pollPendingSatelliteDatagrams(anyInt(), any());
+            waitForIntegerConsumerResult(1);
+            assertEquals(SATELLITE_RESULT_SUCCESS, (int) mIntegerConsumerResult.get(0));
+            mIntegerConsumerResult.clear();
+            verify(mMockSatelliteModemInterface, never()).sendSatelliteDatagram(
+                    any(SatelliteDatagram.class), anyBoolean(), anyBoolean(), any(Message.class));
+
+            moveTimeForward(TIMEOUT_DATAGRAM_DELAY_IN_DEMO_MODE);
+            processAllMessages();
+            verify(mMockDatagramController, times(3)).pushDemoModeDatagram(
+                    anyInt(), any(SatelliteDatagram.class));
+            verify(mMockDatagramController, times(3)).pollPendingSatelliteDatagrams(anyInt(),
+                    any());
+        }
 
         mDatagramDispatcherUT.setDemoMode(false);
         mDatagramDispatcherUT.setDeviceAlignedWithSatellite(false);
diff --git a/tests/telephonytests/src/com/android/internal/telephony/satellite/DatagramReceiverTest.java b/tests/telephonytests/src/com/android/internal/telephony/satellite/DatagramReceiverTest.java
index 79d3657..4d3acb4 100644
--- a/tests/telephonytests/src/com/android/internal/telephony/satellite/DatagramReceiverTest.java
+++ b/tests/telephonytests/src/com/android/internal/telephony/satellite/DatagramReceiverTest.java
@@ -25,6 +25,7 @@
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.anyInt;
 import static org.mockito.Mockito.clearInvocations;
@@ -170,7 +171,7 @@
         }).when(mMockSatelliteModemInterface).pollPendingSatelliteDatagrams(any(Message.class));
         doReturn(true).when(mMockDatagramController)
                 .needsWaitingForSatelliteConnected(eq(SatelliteManager.DATAGRAM_TYPE_UNKNOWN));
-        when(mMockDatagramController.getDatagramWaitTimeForConnectedState())
+        when(mMockDatagramController.getDatagramWaitTimeForConnectedState(anyBoolean()))
                 .thenReturn(TEST_DATAGRAM_WAIT_FOR_CONNECTED_STATE_TIMEOUT_MILLIS);
         mResultListener.clear();
 
@@ -181,7 +182,7 @@
         mInOrder.verify(mMockDatagramController).updateReceiveStatus(eq(SUB_ID),
                 eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT), eq(0),
                 eq(SatelliteManager.SATELLITE_RESULT_SUCCESS));
-        mInOrder.verify(mMockDatagramController).getDatagramWaitTimeForConnectedState();
+        mInOrder.verify(mMockDatagramController).getDatagramWaitTimeForConnectedState(eq(false));
         verifyZeroInteractions(mMockSatelliteModemInterface);
         assertTrue(mDatagramReceiverUT.isDatagramWaitForConnectedStateTimerStarted());
 
@@ -210,7 +211,7 @@
         processAllMessages();
         mInOrder.verify(mMockDatagramController)
                 .needsWaitingForSatelliteConnected(eq(SatelliteManager.DATAGRAM_TYPE_UNKNOWN));
-        mInOrder.verify(mMockDatagramController).getDatagramWaitTimeForConnectedState();
+        mInOrder.verify(mMockDatagramController).getDatagramWaitTimeForConnectedState(eq(false));
         verifyZeroInteractions(mMockSatelliteModemInterface);
         assertTrue(mDatagramReceiverUT.isDatagramWaitForConnectedStateTimerStarted());