From afbd80dc6209f2045d533bf18f4f5ee396c8d4ab Mon Sep 17 00:00:00 2001
From: Ashan Thamara <ashanthamara98@gmail.com>
Date: Thu, 19 Mar 2026 23:46:24 +0530
Subject: [PATCH] Fix issue in resolving device registration cache from session
 data store before expiry

---
 .gitignore                                    |  12 +
 .../pom.xml                                   |   6 +-
 .../DeviceRegistrationRequestCacheEntry.java  |  46 +++-
 .../constant/PushDeviceHandlerConstants.java  |   3 +
 .../DeviceRegistrationContextManagerImpl.java |  18 +-
 ...viceRegistrationRequestCacheEntryTest.java | 209 ++++++++++++++++++
 ...iceRegistrationContextManagerImplTest.java | 105 +++++++--
 7 files changed, 370 insertions(+), 29 deletions(-)
 create mode 100644 components/org.wso2.carbon.identity.notification.push.device.handler/src/test/java/org/wso2/carbon/identity/notification/push/device/handler/cache/DeviceRegistrationRequestCacheEntryTest.java

diff --git a/.gitignore b/.gitignore
index 524f096..024afb0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,14 @@
 # Compiled class file
 *.class
+.classpath
+.factorypath
+.settings
+.project
+*.iml
+*.iws
+*.ipr
+.idea
+.DS_Store
 
 # Log file
 *.log
@@ -22,3 +31,6 @@
 # virtual machine crash logs, see http://www.java.com/en/download/help/error_hotspot.xml
 hs_err_pid*
 replay_pid*
+
+# Ignore everything in this directory
+target
diff --git a/components/org.wso2.carbon.identity.notification.push.device.handler/pom.xml b/components/org.wso2.carbon.identity.notification.push.device.handler/pom.xml
index 0b28d88..324c1d2 100644
--- a/components/org.wso2.carbon.identity.notification.push.device.handler/pom.xml
+++ b/components/org.wso2.carbon.identity.notification.push.device.handler/pom.xml
@@ -1,6 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <!--
-  ~ Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
+  ~ Copyright (c) 2025-2026, WSO2 LLC. (http://www.wso2.com).
   ~
   ~ WSO2 LLC. licenses this file to you under the Apache License,
   ~ Version 2.0 (the "License"); you may not use this file except
@@ -153,6 +153,7 @@
                             org.wso2.carbon.identity.notification.sender.tenant.config.dto; version="${identity.event.handler.notification.imp.pkg.version.range}",
                             org.wso2.carbon.identity.notification.sender.tenant.config.exception; version="${identity.event.handler.notification.imp.pkg.version.range}",
                             org.wso2.carbon.identity.application.authentication.framework.store; version="${carbon.identity.framework.imp.pkg.version.range}",
+                            org.wso2.carbon.identity.application.authentication.framework.util; version="${carbon.identity.framework.imp.pkg.version.range}",
                             org.wso2.carbon.utils; version="${carbon.kernel.package.import.version.range}",
                             org.wso2.carbon.utils.multitenancy; version="${carbon.kernel.package.import.version.range}",
                             org.wso2.carbon.database.utils.jdbc;version="${org.wso2.carbon.database.utils.version.range}",
@@ -189,7 +190,8 @@
                         <exclude>**/*Constants*.class</exclude>
                         <exclude>**/model/**</exclude>
                         <exclude>**/internal/**</exclude>
-                        <exclude>**/cache/**</exclude>
+                        <exclude>**/cache/DeviceRegistrationRequestCache.class</exclude>
+                        <exclude>**/cache/DeviceRegistrationRequestCacheKey.class</exclude>
                     </excludes>
                 </configuration>
                 <executions>
diff --git a/components/org.wso2.carbon.identity.notification.push.device.handler/src/main/java/org/wso2/carbon/identity/notification/push/device/handler/cache/DeviceRegistrationRequestCacheEntry.java b/components/org.wso2.carbon.identity.notification.push.device.handler/src/main/java/org/wso2/carbon/identity/notification/push/device/handler/cache/DeviceRegistrationRequestCacheEntry.java
index 2a7f926..d7ea335 100644
--- a/components/org.wso2.carbon.identity.notification.push.device.handler/src/main/java/org/wso2/carbon/identity/notification/push/device/handler/cache/DeviceRegistrationRequestCacheEntry.java
+++ b/components/org.wso2.carbon.identity.notification.push.device.handler/src/main/java/org/wso2/carbon/identity/notification/push/device/handler/cache/DeviceRegistrationRequestCacheEntry.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
+ * Copyright (c) 2025-2026, WSO2 LLC. (http://www.wso2.com).
  *
  * WSO2 LLC. licenses this file to you under the Apache License,
  * Version 2.0 (the "License"); you may not use this file except
@@ -18,16 +18,29 @@
 
 package org.wso2.carbon.identity.notification.push.device.handler.cache;
 
+import org.apache.commons.lang.StringUtils;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.wso2.carbon.identity.application.authentication.framework.util.FrameworkUtils;
 import org.wso2.carbon.identity.core.cache.CacheEntry;
+import org.wso2.carbon.identity.core.util.IdentityUtil;
 import org.wso2.carbon.identity.notification.push.device.handler.model.DeviceRegistrationContext;
 
+import java.util.concurrent.TimeUnit;
+
+import static org.wso2.carbon.identity.notification.push.device.handler.constant.PushDeviceHandlerConstants.DEFAULT_DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD;
+import static org.wso2.carbon.identity.notification.push.device.handler.constant.PushDeviceHandlerConstants.DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD;
+
 /**
  * Cache entry for Device Registration Request.
  */
 public class DeviceRegistrationRequestCacheEntry extends CacheEntry {
 
     private static final long serialVersionUID = -4976259795783529978L;
+    private static final Log LOG = LogFactory.getLog(DeviceRegistrationRequestCacheEntry.class);
+
     private final DeviceRegistrationContext deviceRegistrationContext;
+    private final long expiresAt;
 
     /**
      * Constructor for DeviceRegistrationRequestCacheEntry.
@@ -37,15 +50,44 @@ public class DeviceRegistrationRequestCacheEntry extends CacheEntry {
     public DeviceRegistrationRequestCacheEntry(DeviceRegistrationContext deviceRegistrationContext) {
 
         this.deviceRegistrationContext = deviceRegistrationContext;
+        long validityPeriodInNanos = getValidityPeriodInNanos();
+        this.expiresAt = getCurrentTimeInNanos() + validityPeriodInNanos;
+        // This will be used in the session store to set the cache entry expiry time.
+        this.setValidityPeriod(validityPeriodInNanos);
     }
 
     /**
      * Get the device registration context.
      *
-     * @return DeviceRegistrationContext
+     * @return DeviceRegistrationContext if the cache entry has not expired, null otherwise.
      */
     public DeviceRegistrationContext getDeviceRegistrationContext() {
 
+        if (getCurrentTimeInNanos() > expiresAt) {
+            LOG.debug("Push device registration request cache entry has expired.");
+            return null;
+        }
         return deviceRegistrationContext;
     }
+
+    private long getValidityPeriodInNanos() {
+
+        int expirySeconds = DEFAULT_DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD;
+        String configuredExpiry = IdentityUtil.getProperty(DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD);
+        if (StringUtils.isNotBlank(configuredExpiry)) {
+            try {
+                expirySeconds = Integer.parseInt(configuredExpiry.trim());
+            } catch (NumberFormatException e) {
+                LOG.warn("Invalid value configured for '" + DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD
+                        + "': " + configuredExpiry + ". Using default value: " + expirySeconds + " seconds.", e);
+            }
+        }
+
+        return TimeUnit.SECONDS.toNanos(expirySeconds);
+    }
+
+    private long getCurrentTimeInNanos() {
+
+        return FrameworkUtils.getCurrentStandardNano();
+    }
 }
diff --git a/components/org.wso2.carbon.identity.notification.push.device.handler/src/main/java/org/wso2/carbon/identity/notification/push/device/handler/constant/PushDeviceHandlerConstants.java b/components/org.wso2.carbon.identity.notification.push.device.handler/src/main/java/org/wso2/carbon/identity/notification/push/device/handler/constant/PushDeviceHandlerConstants.java
index 9da1250..d81bdab 100644
--- a/components/org.wso2.carbon.identity.notification.push.device.handler/src/main/java/org/wso2/carbon/identity/notification/push/device/handler/constant/PushDeviceHandlerConstants.java
+++ b/components/org.wso2.carbon.identity.notification.push.device.handler/src/main/java/org/wso2/carbon/identity/notification/push/device/handler/constant/PushDeviceHandlerConstants.java
@@ -27,6 +27,9 @@ public class PushDeviceHandlerConstants {
     public static final String SIGNATURE_ALGORITHM = "RSA";
     public static final String DEVICE_REGISTRATION_REQUEST_CACHE = "PushDeviceRegistrationRequestCache";
     public static final String DEFAULT_PUSH_PROVIDER = "defaultPushProvider";
+    public static final String DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD =
+            "PushAuthenticator.DeviceRegistrationContext.ValidityPeriod";
+    public static final int DEFAULT_DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD = 180;
 
     /**
      * Private constructor to prevent initialization of the class.
diff --git a/components/org.wso2.carbon.identity.notification.push.device.handler/src/main/java/org/wso2/carbon/identity/notification/push/device/handler/impl/DeviceRegistrationContextManagerImpl.java b/components/org.wso2.carbon.identity.notification.push.device.handler/src/main/java/org/wso2/carbon/identity/notification/push/device/handler/impl/DeviceRegistrationContextManagerImpl.java
index b665f0b..e3b3aca 100644
--- a/components/org.wso2.carbon.identity.notification.push.device.handler/src/main/java/org/wso2/carbon/identity/notification/push/device/handler/impl/DeviceRegistrationContextManagerImpl.java
+++ b/components/org.wso2.carbon.identity.notification.push.device.handler/src/main/java/org/wso2/carbon/identity/notification/push/device/handler/impl/DeviceRegistrationContextManagerImpl.java
@@ -19,7 +19,6 @@
 package org.wso2.carbon.identity.notification.push.device.handler.impl;
 
 import org.wso2.carbon.identity.application.authentication.framework.store.SessionDataStore;
-import org.wso2.carbon.identity.core.util.IdentityCacheUtil;
 import org.wso2.carbon.identity.notification.push.device.handler.DeviceRegistrationContextManager;
 import org.wso2.carbon.identity.notification.push.device.handler.cache.DeviceRegistrationRequestCache;
 import org.wso2.carbon.identity.notification.push.device.handler.cache.DeviceRegistrationRequestCacheEntry;
@@ -36,11 +35,10 @@ public class DeviceRegistrationContextManagerImpl implements DeviceRegistrationC
     @Override
     public void storeRegistrationContext(String key, DeviceRegistrationContext context, String tenantDomain) {
 
-        DeviceRegistrationRequestCache.getInstance().addToCache(
-                new DeviceRegistrationRequestCacheKey(key),
-                new DeviceRegistrationRequestCacheEntry(context),
+        DeviceRegistrationRequestCacheEntry cacheEntry = new DeviceRegistrationRequestCacheEntry(context);
+        DeviceRegistrationRequestCache.getInstance().addToCache(new DeviceRegistrationRequestCacheKey(key), cacheEntry,
                 tenantDomain);
-        storeToSessionStore(key, new DeviceRegistrationRequestCacheEntry(context));
+        storeToSessionStore(key, cacheEntry);
     }
 
     @Override
@@ -48,15 +46,11 @@ public DeviceRegistrationContext getContext(String key, String tenantDomain) {
 
         DeviceRegistrationRequestCacheEntry cacheEntry = DeviceRegistrationRequestCache.getInstance().getValueFromCache(
                 new DeviceRegistrationRequestCacheKey(key), tenantDomain);
-        if (cacheEntry != null) {
-            return cacheEntry.getDeviceRegistrationContext();
-        } else {
+        if (cacheEntry == null) {
             cacheEntry = getFromSessionStore(key);
-            if (!IdentityCacheUtil.isCacheEntryExpired(cacheEntry)) {
-                return cacheEntry.getDeviceRegistrationContext();
-            }
         }
-        return null;
+
+        return cacheEntry != null ? cacheEntry.getDeviceRegistrationContext() : null;
     }
 
     @Override
diff --git a/components/org.wso2.carbon.identity.notification.push.device.handler/src/test/java/org/wso2/carbon/identity/notification/push/device/handler/cache/DeviceRegistrationRequestCacheEntryTest.java b/components/org.wso2.carbon.identity.notification.push.device.handler/src/test/java/org/wso2/carbon/identity/notification/push/device/handler/cache/DeviceRegistrationRequestCacheEntryTest.java
new file mode 100644
index 0000000..48ab6d1
--- /dev/null
+++ b/components/org.wso2.carbon.identity.notification.push.device.handler/src/test/java/org/wso2/carbon/identity/notification/push/device/handler/cache/DeviceRegistrationRequestCacheEntryTest.java
@@ -0,0 +1,209 @@
+/*
+ * Copyright (c) 2026, WSO2 LLC. (http://www.wso2.com).
+ *
+ * WSO2 LLC. licenses this file to you under the Apache License,
+ * Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.wso2.carbon.identity.notification.push.device.handler.cache;
+
+import org.mockito.MockedStatic;
+import org.testng.Assert;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+import org.wso2.carbon.identity.application.authentication.framework.util.FrameworkUtils;
+import org.wso2.carbon.identity.core.util.IdentityUtil;
+import org.wso2.carbon.identity.notification.push.device.handler.model.DeviceRegistrationContext;
+
+import java.util.concurrent.TimeUnit;
+
+import static org.mockito.Mockito.mockStatic;
+import static org.wso2.carbon.identity.notification.push.device.handler.constant.PushDeviceHandlerConstants.DEFAULT_DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD;
+import static org.wso2.carbon.identity.notification.push.device.handler.constant.PushDeviceHandlerConstants.DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD;
+
+/**
+ * Unit tests for DeviceRegistrationRequestCacheEntry.
+ */
+public class DeviceRegistrationRequestCacheEntryTest {
+
+    private static final String TEST_CHALLENGE = "test-challenge";
+    private static final String TEST_USERNAME = "test-user";
+    private static final String TEST_TENANT_DOMAIN = "test.com";
+    private static final long DEFAULT_VALIDITY_NANOS =
+            TimeUnit.SECONDS.toNanos(DEFAULT_DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD);
+
+    private MockedStatic<IdentityUtil> mockedIdentityUtil;
+    private MockedStatic<FrameworkUtils> mockedFrameworkUtils;
+
+    @BeforeMethod
+    public void setUp() {
+
+        mockedIdentityUtil = mockStatic(IdentityUtil.class);
+        mockedFrameworkUtils = mockStatic(FrameworkUtils.class);
+    }
+
+    @AfterMethod
+    public void tearDown() {
+
+        mockedIdentityUtil.close();
+        mockedFrameworkUtils.close();
+    }
+
+    @Test
+    public void testGetDeviceRegistrationContext_ReturnsContext_WhenNotExpired() {
+
+        DeviceRegistrationContext context =
+                new DeviceRegistrationContext(TEST_CHALLENGE, TEST_USERNAME, TEST_TENANT_DOMAIN, false);
+
+        mockedIdentityUtil.when(() -> IdentityUtil.getProperty(DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD))
+                .thenReturn(null);
+        mockedFrameworkUtils.when(FrameworkUtils::getCurrentStandardNano).thenCallRealMethod();
+
+        DeviceRegistrationRequestCacheEntry entry = new DeviceRegistrationRequestCacheEntry(context);
+        DeviceRegistrationContext result = entry.getDeviceRegistrationContext();
+
+        Assert.assertNotNull(result, "Expected context to be returned when not expired.");
+        Assert.assertEquals(result, context, "Expected the same context object to be returned.");
+    }
+
+    @Test
+    public void testGetDeviceRegistrationContext_ReturnsNull_WhenExpired() {
+
+        DeviceRegistrationContext context =
+                new DeviceRegistrationContext(TEST_CHALLENGE, TEST_USERNAME, TEST_TENANT_DOMAIN, false);
+
+        mockedIdentityUtil.when(() -> IdentityUtil.getProperty(DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD))
+                .thenReturn(null);
+        // Constructor at time 0, getter past default validity
+        long expiredQueryTime = DEFAULT_VALIDITY_NANOS + 1L;
+        mockedFrameworkUtils.when(FrameworkUtils::getCurrentStandardNano).thenReturn(0L, expiredQueryTime);
+
+        DeviceRegistrationRequestCacheEntry entry = new DeviceRegistrationRequestCacheEntry(context);
+        DeviceRegistrationContext result = entry.getDeviceRegistrationContext();
+
+        Assert.assertNull(result, "Expected null when the cache entry has expired.");
+    }
+
+    @Test
+    public void testGetDeviceRegistrationContext_ReturnsContext_AtExactExpiry() {
+
+        DeviceRegistrationContext context =
+                new DeviceRegistrationContext(TEST_CHALLENGE, TEST_USERNAME, TEST_TENANT_DOMAIN, false);
+
+        mockedIdentityUtil.when(() -> IdentityUtil.getProperty(DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD))
+                .thenReturn(null);
+        mockedFrameworkUtils.when(FrameworkUtils::getCurrentStandardNano).thenReturn(0L, DEFAULT_VALIDITY_NANOS);
+
+        DeviceRegistrationRequestCacheEntry entry = new DeviceRegistrationRequestCacheEntry(context);
+        DeviceRegistrationContext result = entry.getDeviceRegistrationContext();
+
+        Assert.assertNotNull(result,
+                "Expected context at exact expiry boundary (expiry check uses strict > not >=).");
+    }
+
+    @Test
+    public void testGetDeviceRegistrationContext_ReturnsNull_WhenContextIsNull() {
+
+        mockedIdentityUtil.when(() -> IdentityUtil.getProperty(DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD))
+                .thenReturn(null);
+        mockedFrameworkUtils.when(FrameworkUtils::getCurrentStandardNano).thenReturn(0L, 0L);
+
+        DeviceRegistrationRequestCacheEntry entry = new DeviceRegistrationRequestCacheEntry(null);
+        DeviceRegistrationContext result = entry.getDeviceRegistrationContext();
+
+        Assert.assertNull(result, "Expected null when the stored context itself is null.");
+    }
+
+    @Test
+    public void testValidityPeriod_UsesDefault_WhenConfigIsNull() {
+
+        DeviceRegistrationContext context =
+                new DeviceRegistrationContext(TEST_CHALLENGE, TEST_USERNAME, TEST_TENANT_DOMAIN, false);
+
+        mockedIdentityUtil.when(() -> IdentityUtil.getProperty(DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD))
+                .thenReturn(null);
+        // Constructor: 0L | 1st getter: at exact boundary (still valid) | 2nd getter: past boundary (expired)
+        mockedFrameworkUtils.when(FrameworkUtils::getCurrentStandardNano)
+                .thenReturn(0L, DEFAULT_VALIDITY_NANOS, DEFAULT_VALIDITY_NANOS + 1L);
+
+        DeviceRegistrationRequestCacheEntry entry = new DeviceRegistrationRequestCacheEntry(context);
+
+        Assert.assertNotNull(entry.getDeviceRegistrationContext(),
+                "Context should be valid at the exact default expiry boundary.");
+        Assert.assertNull(entry.getDeviceRegistrationContext(),
+                "Context should be null after the default 180s validity period.");
+    }
+
+    @Test
+    public void testValidityPeriod_UsesDefault_WhenConfigIsBlank() {
+
+        DeviceRegistrationContext context =
+                new DeviceRegistrationContext(TEST_CHALLENGE, TEST_USERNAME, TEST_TENANT_DOMAIN, false);
+
+        mockedIdentityUtil.when(() -> IdentityUtil.getProperty(DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD))
+                .thenReturn("   ");
+        // Constructor: 0L | 1st getter: at exact boundary (still valid) | 2nd getter: past boundary (expired)
+        mockedFrameworkUtils.when(FrameworkUtils::getCurrentStandardNano)
+                .thenReturn(0L, DEFAULT_VALIDITY_NANOS, DEFAULT_VALIDITY_NANOS + 1L);
+
+        DeviceRegistrationRequestCacheEntry entry = new DeviceRegistrationRequestCacheEntry(context);
+
+        Assert.assertNotNull(entry.getDeviceRegistrationContext(),
+                "Context should be valid at the exact default expiry boundary when config is blank.");
+        Assert.assertNull(entry.getDeviceRegistrationContext(),
+                "Context should be null after default validity when config is blank.");
+    }
+
+    @Test
+    public void testValidityPeriod_UsesCustomValue_WhenConfigIsValid() {
+
+        DeviceRegistrationContext context =
+                new DeviceRegistrationContext(TEST_CHALLENGE, TEST_USERNAME, TEST_TENANT_DOMAIN, false);
+        int customSeconds = 300;
+        long customValidityNanos = TimeUnit.SECONDS.toNanos(customSeconds);
+
+        mockedIdentityUtil.when(() -> IdentityUtil.getProperty(DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD))
+                .thenReturn(String.valueOf(customSeconds));
+        // Constructor: 0L
+        // 1st getter: past default 180s but within custom 300s (still valid)
+        // 2nd getter: past custom 300s (expired)
+        mockedFrameworkUtils.when(FrameworkUtils::getCurrentStandardNano)
+                .thenReturn(0L, DEFAULT_VALIDITY_NANOS + 1L, customValidityNanos + 1L);
+
+        DeviceRegistrationRequestCacheEntry entry = new DeviceRegistrationRequestCacheEntry(context);
+
+        Assert.assertNotNull(entry.getDeviceRegistrationContext(),
+                "Context should be valid past default 180s when custom 300s period is configured.");
+        Assert.assertNull(entry.getDeviceRegistrationContext(),
+                "Context should be null after the custom 300s validity period.");
+    }
+
+    @Test
+    public void testValidityPeriod_UsesDefault_WhenConfigIsInvalid() {
+
+        DeviceRegistrationContext context =
+                new DeviceRegistrationContext(TEST_CHALLENGE, TEST_USERNAME, TEST_TENANT_DOMAIN, false);
+
+        mockedIdentityUtil.when(() -> IdentityUtil.getProperty(DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD))
+                .thenReturn("notAnInt");
+        mockedFrameworkUtils.when(FrameworkUtils::getCurrentStandardNano).thenCallRealMethod();
+
+        DeviceRegistrationRequestCacheEntry entry = new DeviceRegistrationRequestCacheEntry(context);
+
+        Assert.assertNotNull(entry.getDeviceRegistrationContext(),
+                "Context should be valid at exact default expiry when config value causes NumberFormatException.");
+        Assert.assertEquals(entry.getValidityPeriod(), DEFAULT_VALIDITY_NANOS);
+    }
+}
diff --git a/components/org.wso2.carbon.identity.notification.push.device.handler/src/test/java/org/wso2/carbon/identity/notification/push/device/handler/impl/DeviceRegistrationContextManagerImplTest.java b/components/org.wso2.carbon.identity.notification.push.device.handler/src/test/java/org/wso2/carbon/identity/notification/push/device/handler/impl/DeviceRegistrationContextManagerImplTest.java
index 18eb235..aa2f58a 100644
--- a/components/org.wso2.carbon.identity.notification.push.device.handler/src/test/java/org/wso2/carbon/identity/notification/push/device/handler/impl/DeviceRegistrationContextManagerImplTest.java
+++ b/components/org.wso2.carbon.identity.notification.push.device.handler/src/test/java/org/wso2/carbon/identity/notification/push/device/handler/impl/DeviceRegistrationContextManagerImplTest.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
+ * Copyright (c) 2025-2026, WSO2 LLC. (http://www.wso2.com).
  *
  * WSO2 LLC. licenses this file to you under the Apache License,
  * Version 2.0 (the "License"); you may not use this file except
@@ -25,21 +25,37 @@
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 import org.wso2.carbon.identity.application.authentication.framework.store.SessionDataStore;
+import org.wso2.carbon.identity.application.authentication.framework.util.FrameworkUtils;
+import org.wso2.carbon.identity.core.util.IdentityUtil;
 import org.wso2.carbon.identity.notification.push.device.handler.cache.DeviceRegistrationRequestCache;
 import org.wso2.carbon.identity.notification.push.device.handler.cache.DeviceRegistrationRequestCacheEntry;
 import org.wso2.carbon.identity.notification.push.device.handler.cache.DeviceRegistrationRequestCacheKey;
 import org.wso2.carbon.identity.notification.push.device.handler.model.DeviceRegistrationContext;
 
+import java.util.concurrent.TimeUnit;
+
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.when;
+import static org.wso2.carbon.identity.notification.push.device.handler.constant.PushDeviceHandlerConstants.DEFAULT_DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD;
+import static org.wso2.carbon.identity.notification.push.device.handler.constant.PushDeviceHandlerConstants.DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD;
 
 /**
  * Unit tests for DeviceRegistrationContextManagerImpl.
  */
 public class DeviceRegistrationContextManagerImplTest {
 
+    private static final String TEST_ID = "id";
+    private static final String TEST_USER = "user";
+    private static final String TEST_TENANT = "tenant";
+    private static final String TEST_KEY = "testKey";
+    private static final String TEST_TENANT_DOMAIN = "carbon.super";
+    private static final String DEVICE_REGISTRATION_REQUEST_CACHE = "DEVICE_REGISTRATION_REQUEST_CACHE";
+    private static final long DEFAULT_VALIDITY_NANOS =
+            TimeUnit.SECONDS.toNanos(DEFAULT_DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD);
+
+
     private DeviceRegistrationContextManagerImpl deviceRegistrationContextManager;
 
     @BeforeMethod
@@ -51,9 +67,9 @@ public void setUp() {
     @Test
     public void testStoreRegistrationContext() {
 
-        DeviceRegistrationContext context = new DeviceRegistrationContext("id", "user", "tenant", false);
-        String key = "testKey";
-        String tenantDomain = "carbon.super";
+        DeviceRegistrationContext context = new DeviceRegistrationContext(TEST_ID, TEST_USER, TEST_TENANT, false);
+        String key = TEST_KEY;
+        String tenantDomain = TEST_TENANT_DOMAIN;
 
         try (MockedStatic<DeviceRegistrationRequestCache> mockedCache
                      = Mockito.mockStatic(DeviceRegistrationRequestCache.class);
@@ -88,9 +104,9 @@ public void testStoreRegistrationContext() {
     @Test
     public void testGetContext() {
 
-        DeviceRegistrationContext context = new DeviceRegistrationContext("id", "user", "tenant", false);
-        String key = "testKey";
-        String tenantDomain = "carbon.super";
+        DeviceRegistrationContext context = new DeviceRegistrationContext(TEST_ID, TEST_USER, TEST_TENANT, false);
+        String key = TEST_KEY;
+        String tenantDomain = TEST_TENANT_DOMAIN;
 
         try (MockedStatic<DeviceRegistrationRequestCache> mockedCache
                      = Mockito.mockStatic(DeviceRegistrationRequestCache.class);
@@ -116,8 +132,8 @@ public void testGetContext() {
     @Test
     public void testGetContextWithNull() {
 
-        String key = "testKey";
-        String tenantDomain = "carbon.super";
+        String key = TEST_KEY;
+        String tenantDomain = TEST_TENANT_DOMAIN;
 
         try (MockedStatic<DeviceRegistrationRequestCache> mockedCache
                      = Mockito.mockStatic(DeviceRegistrationRequestCache.class);
@@ -142,12 +158,12 @@ public void testGetContextWithNull() {
     @Test
     public void testClearContext() {
 
-        String key = "testKey";
-        String tenantDomain = "carbon.super";
+        String key = TEST_KEY;
+        String tenantDomain = TEST_TENANT_DOMAIN;
 
         try (MockedStatic<DeviceRegistrationRequestCache> mockedCache =
                      Mockito.mockStatic(DeviceRegistrationRequestCache.class);
-            MockedStatic<SessionDataStore> mockedSessionDataStore = Mockito.mockStatic(SessionDataStore.class)
+             MockedStatic<SessionDataStore> mockedSessionDataStore = Mockito.mockStatic(SessionDataStore.class)
         ) {
 
             DeviceRegistrationRequestCache cache = Mockito.mock(DeviceRegistrationRequestCache.class);
@@ -155,11 +171,74 @@ public void testClearContext() {
 
             SessionDataStore sessionDataStore = Mockito.mock(SessionDataStore.class);
             mockedSessionDataStore.when(SessionDataStore::getInstance).thenReturn(sessionDataStore);
-            doNothing().when(sessionDataStore).clearSessionData(key, "DEVICE_REGISTRATION_REQUEST_CACHE");
+            doNothing().when(sessionDataStore).clearSessionData(key, DEVICE_REGISTRATION_REQUEST_CACHE);
 
             deviceRegistrationContextManager.clearContext(key, tenantDomain);
 
             Mockito.verify(cache).clearCacheEntry(new DeviceRegistrationRequestCacheKey(key), tenantDomain);
         }
     }
+
+    /**
+     * Test getContext returns null when both the primary cache and the session store miss (return null).
+     */
+    @Test
+    public void testGetContext_ReturnNullWhenBothCacheAndSessionStoreMiss() {
+
+        try (MockedStatic<DeviceRegistrationRequestCache> mockedCache =
+                     Mockito.mockStatic(DeviceRegistrationRequestCache.class);
+             MockedStatic<SessionDataStore> mockedSessionDataStore = Mockito.mockStatic(SessionDataStore.class)
+        ) {
+            DeviceRegistrationRequestCache cache = Mockito.mock(DeviceRegistrationRequestCache.class);
+            mockedCache.when(DeviceRegistrationRequestCache::getInstance).thenReturn(cache);
+            when(cache.getValueFromCache(any(DeviceRegistrationRequestCacheKey.class), anyString()))
+                    .thenReturn(null);
+
+            SessionDataStore sessionDataStore = Mockito.mock(SessionDataStore.class);
+            mockedSessionDataStore.when(SessionDataStore::getInstance).thenReturn(sessionDataStore);
+            when(sessionDataStore.getSessionData(anyString(), anyString())).thenReturn(null);
+
+            DeviceRegistrationContext result = deviceRegistrationContextManager.getContext(TEST_KEY,
+                    TEST_TENANT_DOMAIN);
+
+            Assert.assertNull(result, "Expected null when both cache and session store return null.");
+        }
+    }
+
+    /**
+     * Test getContext returns null when the primary cache misses and the session store entry is expired.
+     */
+    @Test
+    public void testGetContext_ReturnNullWhenSessionStoreEntryIsExpired() {
+
+        DeviceRegistrationContext context = new DeviceRegistrationContext(TEST_ID, TEST_USER, TEST_TENANT, false);
+
+        try (MockedStatic<DeviceRegistrationRequestCache> mockedCache =
+                     Mockito.mockStatic(DeviceRegistrationRequestCache.class);
+             MockedStatic<SessionDataStore> mockedSessionDataStore = Mockito.mockStatic(SessionDataStore.class);
+             MockedStatic<FrameworkUtils> mockedFrameworkUtils = Mockito.mockStatic(FrameworkUtils.class);
+             MockedStatic<IdentityUtil> mockedIdentityUtil = Mockito.mockStatic(IdentityUtil.class)
+        ) {
+            // Primary cache misses.
+            DeviceRegistrationRequestCache cache = Mockito.mock(DeviceRegistrationRequestCache.class);
+            mockedCache.when(DeviceRegistrationRequestCache::getInstance).thenReturn(cache);
+            when(cache.getValueFromCache(any(DeviceRegistrationRequestCacheKey.class), anyString())).thenReturn(null);
+
+            mockedIdentityUtil.when(() -> IdentityUtil.getProperty(DEVICE_REGISTRATION_CONTEXT_VALIDITY_PERIOD))
+                    .thenReturn(null);
+            // Control time: 0L for construction, past validity for getter (expired).
+            mockedFrameworkUtils.when(FrameworkUtils::getCurrentStandardNano)
+                    .thenReturn(0L, DEFAULT_VALIDITY_NANOS + 1L);
+            DeviceRegistrationRequestCacheEntry sessionEntry = new DeviceRegistrationRequestCacheEntry(context);
+
+            SessionDataStore sessionDataStore = Mockito.mock(SessionDataStore.class);
+            mockedSessionDataStore.when(SessionDataStore::getInstance).thenReturn(sessionDataStore);
+            when(sessionDataStore.getSessionData(anyString(), anyString())).thenReturn(sessionEntry);
+
+            DeviceRegistrationContext result = deviceRegistrationContextManager.getContext(TEST_KEY,
+                    TEST_TENANT_DOMAIN);
+
+            Assert.assertNull(result, "Expected null when session store entry is expired.");
+        }
+    }
 }