diff --git a/tests/test_warm_reboot.py b/tests/test_warm_reboot.py
index eec6be98e7fc..e5f2de6ed608 100644
--- a/tests/test_warm_reboot.py
+++ b/tests/test_warm_reboot.py
@@ -5,12 +5,14 @@
 import json
 
 # Get restore count of all processes supporting warm restart
-def swss_get_RestoreCount(state_db):
+def swss_get_RestoreCount(dvs, state_db):
     restore_count = {}
     warmtbl = swsscommon.Table(state_db, swsscommon.STATE_WARM_RESTART_TABLE_NAME)
     keys = warmtbl.getKeys()
     assert  len(keys) !=  0
     for key in keys:
+        if key not in dvs.swssd:
+            continue
         (status, fvs) = warmtbl.get(key)
         assert status == True
         for fv in fvs:
@@ -20,12 +22,14 @@ def swss_get_RestoreCount(state_db):
     return restore_count
 
 # function to check the restore count incremented by 1 for all processes supporting warm restart
-def swss_check_RestoreCount(state_db, restore_count):
+def swss_check_RestoreCount(dvs, state_db, restore_count):
     warmtbl = swsscommon.Table(state_db, swsscommon.STATE_WARM_RESTART_TABLE_NAME)
     keys = warmtbl.getKeys()
     print(keys)
     assert  len(keys) > 0
     for key in keys:
+        if key not in dvs.swssd:
+            continue
         (status, fvs) = warmtbl.get(key)
         assert status == True
         for fv in fvs:
@@ -144,7 +148,7 @@ def test_PortSyncdWarmRestart(dvs, testlog):
     (status, fvs) = neighTbl.get("Ethernet20:11.0.0.10")
     assert status == True
 
-    restore_count = swss_get_RestoreCount(state_db)
+    restore_count = swss_get_RestoreCount(dvs, state_db)
 
     # restart portsyncd
     dvs.runcmd(['sh', '-c', 'pkill -x portsyncd'])
@@ -255,7 +259,7 @@ def test_VlanMgrdWarmRestart(dvs, testlog):
     (exitcode, bv_before) = dvs.runcmd("bridge vlan")
     print(bv_before)
 
-    restore_count = swss_get_RestoreCount(state_db)
+    restore_count = swss_get_RestoreCount(dvs, state_db)
 
     dvs.runcmd(['sh', '-c', 'pkill -x vlanmgrd'])
 
@@ -376,7 +380,7 @@ def test_swss_neighbor_syncup(dvs, testlog):
     #
 
     # get restore_count
-    restore_count = swss_get_RestoreCount(state_db)
+    restore_count = swss_get_RestoreCount(dvs, state_db)
 
     # stop neighsyncd and sairedis.rec
     stop_neighsyncd(dvs)
@@ -425,7 +429,7 @@ def test_swss_neighbor_syncup(dvs, testlog):
     #       just that if we want to add the same neighbor again, use "change" instead of "add"
 
     # get restore_count
-    restore_count = swss_get_RestoreCount(state_db)
+    restore_count = swss_get_RestoreCount(dvs, state_db)
 
     # stop neighsyncd
     stop_neighsyncd(dvs)
@@ -496,7 +500,7 @@ def test_swss_neighbor_syncup(dvs, testlog):
     # Check the timer is not retrieved from configDB since it is not configured
 
     # get restore_count
-    restore_count = swss_get_RestoreCount(state_db)
+    restore_count = swss_get_RestoreCount(dvs, state_db)
 
     # stop neighsyncd
     stop_neighsyncd(dvs)
@@ -571,7 +575,7 @@ def test_swss_neighbor_syncup(dvs, testlog):
     dvs.runcmd("config warm_restart neighsyncd_timer {}".format(timer_value))
 
     # get restore_count
-    restore_count = swss_get_RestoreCount(state_db)
+    restore_count = swss_get_RestoreCount(dvs, state_db)
 
     # stop neighsyncd
     stop_neighsyncd(dvs)
@@ -719,7 +723,7 @@ def test_swss_port_state_syncup(dvs, testlog):
 
     tbl = swsscommon.Table(appl_db, swsscommon.APP_PORT_TABLE_NAME)
 
-    restore_count = swss_get_RestoreCount(state_db)
+    restore_count = swss_get_RestoreCount(dvs, state_db)
 
     # update port admin state
     dvs.runcmd("ifconfig Ethernet0 10.0.0.0/31 up")
@@ -766,7 +770,7 @@ def test_swss_port_state_syncup(dvs, testlog):
     dvs.start_swss()
     time.sleep(10)
 
-    swss_check_RestoreCount(state_db, restore_count)
+    swss_check_RestoreCount(dvs, state_db, restore_count)
 
     for i in [0, 1, 2]:
         (status, fvs) = tbl.get("Ethernet%d" % (i * 4))