Merge "Restart P2P_SMS inactivity timer when device receives SMS" into main
diff --git a/src/java/com/android/internal/telephony/InboundSmsHandler.java b/src/java/com/android/internal/telephony/InboundSmsHandler.java
index ca03f5d..0106516 100644
--- a/src/java/com/android/internal/telephony/InboundSmsHandler.java
+++ b/src/java/com/android/internal/telephony/InboundSmsHandler.java
@@ -77,6 +77,7 @@
 import com.android.internal.telephony.flags.FeatureFlags;
 import com.android.internal.telephony.flags.Flags;
 import com.android.internal.telephony.metrics.TelephonyMetrics;
+import com.android.internal.telephony.satellite.SatelliteController;
 import com.android.internal.telephony.satellite.metrics.CarrierRoamingSatelliteSessionStats;
 import com.android.internal.telephony.util.NotificationChannelController;
 import com.android.internal.telephony.util.TelephonyUtils;
@@ -687,6 +688,17 @@
             result = RESULT_SMS_DISPATCH_FAILURE;
         }
 
+        if (mFeatureFlags.carrierRoamingNbIotNtn()) {
+            if (result == Intents.RESULT_SMS_HANDLED) {
+                SatelliteController satelliteController = SatelliteController.getInstance();
+                if (satelliteController == null) {
+                    log("SatelliteController is not initialized");
+                    return;
+                }
+                satelliteController.onSmsReceived(mPhone.getSubId());
+            }
+        }
+
         // RESULT_OK means that the SMS will be acknowledged by special handling,
         // e.g. for SMS-PP data download. Any other result, we should ack here.
         if (result != Activity.RESULT_OK) {
diff --git a/src/java/com/android/internal/telephony/satellite/DatagramController.java b/src/java/com/android/internal/telephony/satellite/DatagramController.java
index 5cd2f99..a1e4bf5 100644
--- a/src/java/com/android/internal/telephony/satellite/DatagramController.java
+++ b/src/java/com/android/internal/telephony/satellite/DatagramController.java
@@ -282,7 +282,7 @@
             mSendDatagramTransferState = datagramTransferState;
             mSendPendingCount = sendPendingCount;
             mSendErrorCode = errorCode;
-            notifyDatagramTransferStateChangedToSessionController();
+            notifyDatagramTransferStateChangedToSessionController(mDatagramType);
             mPointingAppController.updateSendDatagramTransferState(mSendSubId, mDatagramType,
                     mSendDatagramTransferState, mSendPendingCount, mSendErrorCode);
             retryPollPendingDatagramsInDemoMode();
@@ -311,21 +311,23 @@
      * @param receivePendingCount The number of datagrams that are currently pending to be received.
      * @param errorCode If datagram transfer failed, the reason for failure.
      */
-    public void updateReceiveStatus(int subId,
+    public void updateReceiveStatus(int subId, @SatelliteManager.DatagramType int datagramType,
             @SatelliteManager.SatelliteDatagramTransferState int datagramTransferState,
             int receivePendingCount, int errorCode) {
         synchronized (mLock) {
             plogd("updateReceiveStatus"
                     + " subId: " + subId
+                    + " datagramType: " + datagramType
                     + " datagramTransferState: " + datagramTransferState
                     + " receivePendingCount: " + receivePendingCount + " errorCode: " + errorCode);
 
             mReceiveSubId = subId;
+            mDatagramType = datagramType;
             mReceiveDatagramTransferState = datagramTransferState;
             mReceivePendingCount = receivePendingCount;
             mReceiveErrorCode = errorCode;
 
-            notifyDatagramTransferStateChangedToSessionController();
+            notifyDatagramTransferStateChangedToSessionController(mDatagramType);
             mPointingAppController.updateReceiveDatagramTransferState(mReceiveSubId,
                     mReceiveDatagramTransferState, mReceivePendingCount, mReceiveErrorCode);
             retryPollPendingDatagramsInDemoMode();
@@ -365,6 +367,24 @@
     }
 
     /**
+     * Notify SMS received.
+     *
+     * @param subId The subId of the subscription used to receive SMS
+     */
+    public void onSmsReceived(int subId) {
+        // To keep exist notification flow, need to call with each state.
+        updateReceiveStatus(subId, SatelliteManager.DATAGRAM_TYPE_SMS,
+                SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING,
+                getReceivePendingCount(), SatelliteManager.SATELLITE_RESULT_SUCCESS);
+        updateReceiveStatus(subId, SatelliteManager.DATAGRAM_TYPE_SMS,
+                SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_SUCCESS,
+                getReceivePendingCount(), SatelliteManager.SATELLITE_RESULT_SUCCESS);
+        updateReceiveStatus(subId, SatelliteManager.DATAGRAM_TYPE_SMS,
+                SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                getReceivePendingCount(), SatelliteManager.SATELLITE_RESULT_SUCCESS);
+    }
+
+    /**
      * Set whether the device is aligned with the satellite.
      */
     @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
@@ -490,16 +510,6 @@
         }
     }
 
-    /** Return the datagram type indicating that the message to be sent or received.
-     * {@link SatelliteManager.DatagramType}
-     */
-    @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
-    public int getDatagramType() {
-        synchronized (mLock) {
-            return mDatagramType;
-        }
-    }
-
     /**
      * This API can be used by only CTS to timeout durations used by DatagramController module.
      *
@@ -575,14 +585,16 @@
         return (DEBUG || SystemProperties.getBoolean(ALLOW_MOCK_MODEM_PROPERTY, false));
     }
 
-    private void notifyDatagramTransferStateChangedToSessionController() {
+    private void notifyDatagramTransferStateChangedToSessionController(int datagramType) {
         SatelliteSessionController sessionController = SatelliteSessionController.getInstance();
         if (sessionController == null) {
             ploge("notifyDatagramTransferStateChangeToSessionController: SatelliteSessionController"
                     + " is not initialized yet");
         } else {
-            sessionController.onDatagramTransferStateChanged(
-                    mSendDatagramTransferState, mReceiveDatagramTransferState);
+            synchronized (mLock) {
+                sessionController.onDatagramTransferStateChanged(
+                        mSendDatagramTransferState, mReceiveDatagramTransferState, datagramType);
+            }
         }
     }
 
diff --git a/src/java/com/android/internal/telephony/satellite/DatagramReceiver.java b/src/java/com/android/internal/telephony/satellite/DatagramReceiver.java
index 5a89c40..1120f58 100644
--- a/src/java/com/android/internal/telephony/satellite/DatagramReceiver.java
+++ b/src/java/com/android/internal/telephony/satellite/DatagramReceiver.java
@@ -352,10 +352,12 @@
 
                     if (pendingCount <= 0 && satelliteDatagram == null) {
                         sInstance.mDatagramController.updateReceiveStatus(mSubId,
+                                SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE,
                                 SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_NONE,
                                 pendingCount, SatelliteManager.SATELLITE_RESULT_SUCCESS);
                     } else if (satelliteDatagram != null) {
                         sInstance.mDatagramController.updateReceiveStatus(mSubId,
+                                SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE,
                                 SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_SUCCESS,
                                 pendingCount, SatelliteManager.SATELLITE_RESULT_SUCCESS);
 
@@ -376,6 +378,7 @@
 
                     if (pendingCount <= 0) {
                         sInstance.mDatagramController.updateReceiveStatus(mSubId,
+                                SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE,
                                 SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
                                 pendingCount, SatelliteManager.SATELLITE_RESULT_SUCCESS);
                     } else {
@@ -473,10 +476,12 @@
                 plogd("EVENT_POLL_PENDING_SATELLITE_DATAGRAMS_DONE error: " + error);
                 if (error != SatelliteManager.SATELLITE_RESULT_SUCCESS) {
                     mDatagramController.updateReceiveStatus(request.subId,
+                            SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE,
                             SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_FAILED,
                             mDatagramController.getReceivePendingCount(), error);
 
                     mDatagramController.updateReceiveStatus(request.subId,
+                            SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE,
                             SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
                             mDatagramController.getReceivePendingCount(),
                             SatelliteManager.SATELLITE_RESULT_SUCCESS);
@@ -615,6 +620,7 @@
                 mPendingPollSatelliteDatagramsRequest = new DatagramReceiverHandlerRequest(
                         callback, SatelliteServiceUtils.getPhone(), subId);
                 mDatagramController.updateReceiveStatus(subId,
+                        SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE,
                         SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT,
                         mDatagramController.getReceivePendingCount(),
                         SatelliteManager.SATELLITE_RESULT_SUCCESS);
@@ -624,6 +630,7 @@
         }
 
         mDatagramController.updateReceiveStatus(subId,
+                SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE,
                 SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING,
                 mDatagramController.getReceivePendingCount(),
                 SatelliteManager.SATELLITE_RESULT_SUCCESS);
@@ -699,11 +706,13 @@
         int subId = SatelliteController.getInstance().getHighestPrioritySubscrption();
         if (mDatagramController.isReceivingDatagrams()) {
             mDatagramController.updateReceiveStatus(subId,
+                    SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE,
                     SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_FAILED,
                     mDatagramController.getReceivePendingCount(),
                     SatelliteManager.SATELLITE_RESULT_REQUEST_ABORTED);
         }
         mDatagramController.updateReceiveStatus(subId,
+                SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE,
                 SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE, 0,
                 SatelliteManager.SATELLITE_RESULT_SUCCESS);
         cleanupDemoModeResources();
@@ -857,11 +866,13 @@
 
             plogw("Timed out to wait for satellite connected before polling datagrams");
             mDatagramController.updateReceiveStatus(mPendingPollSatelliteDatagramsRequest.subId,
+                    SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE,
                     SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_FAILED,
                     mDatagramController.getReceivePendingCount(),
                     SatelliteManager.SATELLITE_RESULT_NOT_REACHABLE);
 
             mDatagramController.updateReceiveStatus(mPendingPollSatelliteDatagramsRequest.subId,
+                    SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE,
                     SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
                     mDatagramController.getReceivePendingCount(),
                     SatelliteManager.SATELLITE_RESULT_SUCCESS);
diff --git a/src/java/com/android/internal/telephony/satellite/SatelliteController.java b/src/java/com/android/internal/telephony/satellite/SatelliteController.java
index 6466517..69cc08f 100644
--- a/src/java/com/android/internal/telephony/satellite/SatelliteController.java
+++ b/src/java/com/android/internal/telephony/satellite/SatelliteController.java
@@ -3257,6 +3257,36 @@
     }
 
     /**
+     * Notify SMS received.
+     *
+     * @param subId The subId of the subscription used to receive SMS
+     */
+    public void onSmsReceived(int subId) {
+        if (!mFeatureFlags.carrierRoamingNbIotNtn()) {
+            logd("onSmsReceived: carrierRoamingNbIotNtn is disabled");
+            return;
+        }
+
+        if (!isSatelliteEnabled()) {
+            logd("onSmsReceived: satellite is not enabled");
+            return;
+        }
+
+        int satelliteSubId = getHighestPrioritySubscrption();
+        if (subId != satelliteSubId) {
+            logd("onSmsReceived: SMS received " + subId
+                    + ", but not satellite subscription " + satelliteSubId);
+            return;
+        }
+
+        if (mDatagramController != null) {
+            mDatagramController.onSmsReceived(subId);
+        } else {
+            logd("onSmsReceived: DatagramController is not initialized");
+        }
+    }
+
+    /**
      * @return {@code true} if satellite is supported via OEM on the device,
      * {@code  false} otherwise.
      */
diff --git a/src/java/com/android/internal/telephony/satellite/SatelliteSessionController.java b/src/java/com/android/internal/telephony/satellite/SatelliteSessionController.java
index 5b7edd4..55e6b2e 100644
--- a/src/java/com/android/internal/telephony/satellite/SatelliteSessionController.java
+++ b/src/java/com/android/internal/telephony/satellite/SatelliteSessionController.java
@@ -334,9 +334,10 @@
     @VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
     public void onDatagramTransferStateChanged(
             @SatelliteManager.SatelliteDatagramTransferState int sendState,
-            @SatelliteManager.SatelliteDatagramTransferState int receiveState) {
+            @SatelliteManager.SatelliteDatagramTransferState int receiveState,
+            @SatelliteManager.DatagramType int datagramType) {
         sendMessage(EVENT_DATAGRAM_TRANSFER_STATE_CHANGED,
-                new DatagramTransferState(sendState, receiveState));
+                new DatagramTransferState(sendState, receiveState, datagramType));
         if (sendState == SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING) {
             mIsSendingTriggeredDuringTransferringState.set(true);
         }
@@ -627,11 +628,14 @@
     private static class DatagramTransferState {
         @SatelliteManager.SatelliteDatagramTransferState public int sendState;
         @SatelliteManager.SatelliteDatagramTransferState public int receiveState;
+        @SatelliteManager.DatagramType public int datagramType;
 
         DatagramTransferState(@SatelliteManager.SatelliteDatagramTransferState int sendState,
-                @SatelliteManager.SatelliteDatagramTransferState int receiveState) {
+                @SatelliteManager.SatelliteDatagramTransferState int receiveState,
+                @SatelliteManager.DatagramType int datagramType) {
             this.sendState = sendState;
             this.receiveState = receiveState;
+            this.datagramType = datagramType;
         }
     }
 
@@ -1292,8 +1296,7 @@
                     || isReceiving(datagramTransferState.receiveState)) {
                 stopNbIotInactivityTimer();
 
-                DatagramController datagramController = DatagramController.getInstance();
-                int datagramType = datagramController.getDatagramType();
+                int datagramType = datagramTransferState.datagramType;
                 if (datagramType == DATAGRAM_TYPE_SOS_MESSAGE) {
                     stopEsosInactivityTimer();
                 } else if (datagramType == DATAGRAM_TYPE_SMS) {
diff --git a/tests/telephonytests/src/com/android/internal/telephony/satellite/DatagramControllerTest.java b/tests/telephonytests/src/com/android/internal/telephony/satellite/DatagramControllerTest.java
index 2961b4d..21731f0 100644
--- a/tests/telephonytests/src/com/android/internal/telephony/satellite/DatagramControllerTest.java
+++ b/tests/telephonytests/src/com/android/internal/telephony/satellite/DatagramControllerTest.java
@@ -92,8 +92,8 @@
         // Move both send and receive to IDLE state
         mDatagramControllerUT.updateSendStatus(SUB_ID, DATAGRAM_TYPE_UNKNOWN,
                 SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE, 0, SATELLITE_RESULT_SUCCESS);
-        mDatagramControllerUT.updateReceiveStatus(SUB_ID, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE, 0,
-                SATELLITE_RESULT_SUCCESS);
+        mDatagramControllerUT.updateReceiveStatus(SUB_ID, DATAGRAM_TYPE_SOS_MESSAGE,
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE, 0, SATELLITE_RESULT_SUCCESS);
         pushDemoModeSosDatagram(DATAGRAM_TYPE_SOS_MESSAGE);
     }
 
@@ -197,7 +197,7 @@
                 errorCode);
 
         verify(mMockSatelliteSessionController)
-                .onDatagramTransferStateChanged(eq(sendState), anyInt());
+                .onDatagramTransferStateChanged(eq(sendState), anyInt(), anyInt());
         verify(mMockPointingAppController).updateSendDatagramTransferState(
                 eq(SUB_ID), eq(datagramType), eq(sendState), eq(sendPendingCount), eq(errorCode));
 
@@ -219,10 +219,10 @@
         int receivePendingCount = 1;
         int errorCode = SATELLITE_RESULT_SUCCESS;
         mDatagramControllerUT.updateReceiveStatus(
-                SUB_ID, receiveState, receivePendingCount, errorCode);
+                SUB_ID, DATAGRAM_TYPE_SOS_MESSAGE, receiveState, receivePendingCount, errorCode);
 
         verify(mMockSatelliteSessionController)
-                .onDatagramTransferStateChanged(anyInt(), eq(receiveState));
+                .onDatagramTransferStateChanged(anyInt(), eq(receiveState), anyInt());
         verify(mMockPointingAppController).updateReceiveDatagramTransferState(
                 eq(SUB_ID), eq(receiveState), eq(receivePendingCount), eq(errorCode));
 
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 e964ced..a1f63d0 100644
--- a/tests/telephonytests/src/com/android/internal/telephony/satellite/DatagramReceiverTest.java
+++ b/tests/telephonytests/src/com/android/internal/telephony/satellite/DatagramReceiverTest.java
@@ -187,6 +187,7 @@
         mInOrder.verify(mMockDatagramController)
                 .needsWaitingForSatelliteConnected(eq(SatelliteManager.DATAGRAM_TYPE_UNKNOWN));
         mInOrder.verify(mMockDatagramController).updateReceiveStatus(eq(SUB_ID),
+                eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                 eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT), eq(0),
                 eq(SatelliteManager.SATELLITE_RESULT_SUCCESS));
         mInOrder.verify(mMockDatagramController).getDatagramWaitTimeForConnectedState(eq(false));
@@ -201,6 +202,7 @@
 
         mInOrder.verify(mMockDatagramController)
                 .updateReceiveStatus(eq(SUB_ID),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING), eq(0),
                         eq(SatelliteManager.SATELLITE_RESULT_SUCCESS));
         verify(mMockSatelliteModemInterface, times(1))
@@ -225,9 +227,11 @@
         moveTimeForward(TEST_DATAGRAM_WAIT_FOR_CONNECTED_STATE_TIMEOUT_MILLIS);
         processAllMessages();
         mInOrder.verify(mMockDatagramController).updateReceiveStatus(eq(SUB_ID),
+                eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                 eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_FAILED), eq(0),
                 eq(SatelliteManager.SATELLITE_RESULT_NOT_REACHABLE));
         mInOrder.verify(mMockDatagramController).updateReceiveStatus(eq(SUB_ID),
+                eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                 eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE), eq(0),
                 eq(SatelliteManager.SATELLITE_RESULT_SUCCESS));
         verifyZeroInteractions(mMockSatelliteModemInterface);
@@ -266,10 +270,12 @@
 
         mInOrder.verify(mMockDatagramController)
                 .updateReceiveStatus(eq(SUB_ID),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING), eq(0),
                         eq(SatelliteManager.SATELLITE_RESULT_SUCCESS));
         mInOrder.verify(mMockDatagramController)
                 .updateReceiveStatus(eq(SUB_ID),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_FAILED),
                         eq(0), eq(SatelliteManager.SATELLITE_RESULT_SERVICE_ERROR));
 
@@ -289,10 +295,12 @@
 
         mInOrder.verify(mMockDatagramController)
                 .updateReceiveStatus(eq(SUB_ID),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_NONE),
                         eq(0), eq(SatelliteManager.SATELLITE_RESULT_SUCCESS));
         mInOrder.verify(mMockDatagramController)
                 .updateReceiveStatus(eq(SUB_ID),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE),
                         eq(0), eq(SatelliteManager.SATELLITE_RESULT_SUCCESS));
 
@@ -312,10 +320,12 @@
 
         mInOrder.verify(mMockDatagramController)
                 .updateReceiveStatus(eq(SUB_ID),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_SUCCESS),
                         eq(0), eq(SatelliteManager.SATELLITE_RESULT_SUCCESS));
         mInOrder.verify(mMockDatagramController)
                 .updateReceiveStatus(eq(SUB_ID),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE),
                         eq(0), eq(SatelliteManager.SATELLITE_RESULT_SUCCESS));
         assertTrue(testSatelliteDatagramCallback.waitForOnSatelliteDatagramReceived());
@@ -339,6 +349,7 @@
 
         mInOrder.verify(mMockDatagramController)
                 .updateReceiveStatus(eq(SUB_ID),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_SUCCESS),
                         eq(10), eq(SatelliteManager.SATELLITE_RESULT_SUCCESS));
         verify(mMockSessionMetricsStats, times(1)).addCountOfSuccessfulIncomingDatagram();
@@ -356,16 +367,19 @@
         verify(mMockDatagramController, times(1)).popDemoModeDatagram();
         verify(mMockDatagramController)
                 .updateReceiveStatus(eq(SUB_ID),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING),
                         anyInt(),
                         eq(SatelliteManager.SATELLITE_RESULT_SUCCESS));
         verify(mMockDatagramController)
                 .updateReceiveStatus(eq(SUB_ID),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_FAILED),
                         anyInt(),
                         eq(SatelliteManager.SATELLITE_RESULT_INVALID_TELEPHONY_STATE));
         verify(mMockDatagramController)
                 .updateReceiveStatus(eq(SUB_ID),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE),
                         anyInt(),
                         eq(SatelliteManager.SATELLITE_RESULT_SUCCESS));
@@ -387,17 +401,20 @@
         verify(mMockDatagramController, never()).popDemoModeDatagram();
         verify(mMockDatagramController)
                 .updateReceiveStatus(eq(SUB_ID),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING),
                         anyInt(),
                         eq(SatelliteManager.SATELLITE_RESULT_SUCCESS));
         processAllFutureMessages();
         verify(mMockDatagramController)
                 .updateReceiveStatus(eq(SUB_ID),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_FAILED),
                         anyInt(),
                         eq(SatelliteManager.SATELLITE_RESULT_NOT_REACHABLE));
         verify(mMockDatagramController)
                 .updateReceiveStatus(eq(SUB_ID),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE),
                         anyInt(),
                         eq(SatelliteManager.SATELLITE_RESULT_SUCCESS));
@@ -449,10 +466,12 @@
 
         mInOrder.verify(mMockDatagramController)
                 .updateReceiveStatus(anyInt(),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_FAILED),
                         eq(10), eq(SatelliteManager.SATELLITE_RESULT_REQUEST_ABORTED));
         mInOrder.verify(mMockDatagramController)
                 .updateReceiveStatus(anyInt(),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE),
                         eq(0), eq(SatelliteManager.SATELLITE_RESULT_SUCCESS));
     }
@@ -468,6 +487,7 @@
 
         mInOrder.verify(mMockDatagramController)
                 .updateReceiveStatus(anyInt(),
+                        eq(SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE),
                         eq(SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE),
                         eq(0), eq(SatelliteManager.SATELLITE_RESULT_SUCCESS));
     }
diff --git a/tests/telephonytests/src/com/android/internal/telephony/satellite/SatelliteSessionControllerTest.java b/tests/telephonytests/src/com/android/internal/telephony/satellite/SatelliteSessionControllerTest.java
index 7845352..d8daa76 100644
--- a/tests/telephonytests/src/com/android/internal/telephony/satellite/SatelliteSessionControllerTest.java
+++ b/tests/telephonytests/src/com/android/internal/telephony/satellite/SatelliteSessionControllerTest.java
@@ -21,6 +21,7 @@
 import static android.telephony.CarrierConfigManager.KEY_SATELLITE_ROAMING_SCREEN_OFF_INACTIVITY_TIMEOUT_SEC_INT;
 import static android.telephony.satellite.SatelliteManager.DATAGRAM_TYPE_SMS;
 import static android.telephony.satellite.SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE;
+import static android.telephony.satellite.SatelliteManager.DATAGRAM_TYPE_UNKNOWN;
 import static android.telephony.satellite.SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE;
 import static android.telephony.satellite.SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_FAILED;
 import static android.telephony.satellite.SatelliteManager.SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_SUCCESS;
@@ -301,8 +302,6 @@
     public void testP2pSmsInactivityTimer() {
         when(mFeatureFlags.carrierRoamingNbIotNtn()).thenReturn(true);
         when(mMockSatelliteController.isSatelliteAttachRequired()).thenReturn(true);
-        when(mMockDatagramController.getDatagramType()).thenReturn(
-                SatelliteManager.DATAGRAM_TYPE_SMS);
 
         when(mMockSatelliteController.getRequestIsEmergency()).thenReturn(false);
         when(mMockSatelliteController.isSatelliteRoamingP2pSmSSupported(
@@ -347,8 +346,7 @@
     public void testEsosInactivityTimer() {
         when(mFeatureFlags.carrierRoamingNbIotNtn()).thenReturn(true);
         when(mMockSatelliteController.isSatelliteAttachRequired()).thenReturn(true);
-        when(mMockDatagramController.getDatagramType()).thenReturn(
-                SatelliteManager.DATAGRAM_TYPE_SOS_MESSAGE);
+
         when(mMockSatelliteController.getRequestIsEmergency()).thenReturn(true);
         when(mMockSatelliteController.isSatelliteEsosSupported(anyInt())).thenReturn(true);
         PersistableBundle bundle = new PersistableBundle();
@@ -423,7 +421,8 @@
         // Notify datagram controller is in WAITING_TO_CONNECT.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // Verify that ESOS, P2P_SMS timer are not started.
@@ -505,7 +504,8 @@
         // Notify datagram controller is in WAITING_TO_CONNECT.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // Verify that ESOS, P2P_SMS timer are not started.
@@ -579,7 +579,8 @@
         // Notify datagram controller is in WAITING_TO_CONNECT.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // Verify that ESOS, P2P_SMS timer are not started.
@@ -645,7 +646,8 @@
         // Notify datagram controller is in WAITING_TO_CONNECT.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // Verify that ESOS, P2P_SMS timer are not started.
@@ -716,7 +718,8 @@
         // Notify datagram controller is in WAITING_TO_CONNECT.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // Verify that ESOS, P2P_SMS timer are not started.
@@ -938,7 +941,8 @@
 
         // Start sending datagrams
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
-                SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to TRANSFERRING state.
@@ -950,7 +954,8 @@
         // Sending datagrams failed
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_SEND_FAILED,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to IDLE state.
@@ -962,7 +967,8 @@
         // Start sending datagrams again
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to TRANSFERRING state.
@@ -974,7 +980,8 @@
         // Sending datagrams is successful and done.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to LISTENING state.
@@ -987,7 +994,8 @@
         // Start receiving datagrams
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to TRANSFERRING state.
@@ -999,7 +1007,8 @@
 
         // Receiving datagrams is successful and done.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to LISTENING state.
@@ -1012,7 +1021,8 @@
         // Start receiving datagrams again
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to TRANSFERRING state.
@@ -1025,7 +1035,8 @@
         // Receiving datagrams failed.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_FAILED);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_FAILED,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to IDLE state.
@@ -1037,7 +1048,8 @@
         // Start receiving datagrams again
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to TRANSFERRING state.
@@ -1048,7 +1060,8 @@
 
         // Receiving datagrams is successful and done.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to LISTENING state.
@@ -1072,7 +1085,8 @@
         // Start receiving datagrams again
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to TRANSFERRING state.
@@ -1084,7 +1098,8 @@
         // Start sending datagrams
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should stay at TRANSFERRING state.
@@ -1095,7 +1110,8 @@
         // Receiving datagrams failed.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_FAILED);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_FAILED,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should stay at TRANSFERRING state instead of moving to IDLE
@@ -1107,7 +1123,8 @@
         // Start receiving datagrams again.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should stay at TRANSFERRING state.
@@ -1118,7 +1135,8 @@
         // Sending datagrams failed.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_SEND_FAILED,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should stay at TRANSFERRING state instead of moving to IDLE
@@ -1194,7 +1212,8 @@
 
         // Start sending datagrams
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
-                SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // The datagram sending event should be ignored.
@@ -1217,7 +1236,8 @@
 
         // Start sending datagrams
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
-                SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to TRANSFERRING state.
@@ -1232,7 +1252,8 @@
         // Sending datagrams failed
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_SEND_FAILED,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to CONNECTED state.
@@ -1247,7 +1268,8 @@
         // Start sending datagrams again
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to TRANSFERRING state.
@@ -1262,7 +1284,8 @@
         // Sending datagrams is successful and done.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to CONNECTED state.
@@ -1277,7 +1300,8 @@
         // Start receiving datagrams
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVING,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to TRANSFERRING state.
@@ -1291,7 +1315,8 @@
 
         // Receiving datagrams is successful and done.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to CONNECTED state.
@@ -1306,7 +1331,8 @@
         // Start receiving datagrams
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_SUCCESS);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_RECEIVE_SUCCESS,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to TRANSFERRING state.
@@ -1320,7 +1346,8 @@
 
         // Receiving datagrams is successful and done.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to CONNECTED state.
@@ -1348,7 +1375,8 @@
         // Start sending datagrams
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to NOT_CONNECTED state.
@@ -1446,7 +1474,8 @@
         // Start sending datagrams
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should stay at IDLE state because it failed to disable
@@ -1502,7 +1531,8 @@
         // Start sending datagrams and the NB-IOT inactivity timer should be stopped.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_WAITING_TO_CONNECT,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         moveTimeForward(TEST_SATELLITE_TIMEOUT_MILLIS);
         processAllMessages();
 
@@ -1514,10 +1544,12 @@
         // The NB-IOT inactivity timer should be started.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_SEND_FAILED,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
         assertTrue(mTestSatelliteSessionController.isNbIotInactivityTimerStarted());
 
@@ -1775,17 +1807,15 @@
     }
 
     private void sendMessage(@SatelliteManager.DatagramType int datagramType) {
-        when(mMockDatagramController.getDatagramType()).thenReturn(datagramType);
-
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE, datagramType);
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_SEND_SUCCESS,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE, datagramType);
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE, datagramType);
 
         processAllMessages();
     }
@@ -1853,7 +1883,8 @@
         assertEquals(STATE_IDLE, mTestSatelliteSessionController.getCurrentStateName());
         // Start sending datagrams
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
-                SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_SENDING, SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to TRANSFERRING state.
@@ -1868,7 +1899,8 @@
         // Sending datagrams is successful and done.
         mTestSatelliteSessionController.onDatagramTransferStateChanged(
                 SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
-                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE);
+                SATELLITE_DATAGRAM_TRANSFER_STATE_IDLE,
+                DATAGRAM_TYPE_UNKNOWN);
         processAllMessages();
 
         // SatelliteSessionController should move to LISTENING state.