diff --git a/bl31/bl31_main.c b/bl31/bl31_main.c
index 861b391..ff3c53b 100644
--- a/bl31/bl31_main.c
+++ b/bl31/bl31_main.c
@@ -125,7 +125,7 @@
  ******************************************************************************/
 void bl31_set_next_image_type(uint32_t security_state)
 {
-	assert(security_state == NON_SECURE || security_state == SECURE);
+	assert(sec_state_is_valid(security_state));
 	next_image_type = security_state;
 }
 
diff --git a/bl31/context_mgmt.c b/bl31/context_mgmt.c
index 65f1213..4502e5d 100644
--- a/bl31/context_mgmt.c
+++ b/bl31/context_mgmt.c
@@ -71,7 +71,7 @@
  ******************************************************************************/
 void *cm_get_context_by_mpidr(uint64_t mpidr, uint32_t security_state)
 {
-	assert(security_state <= NON_SECURE);
+	assert(sec_state_is_valid(security_state));
 
 	return get_cpu_data_by_mpidr(mpidr, cpu_context[security_state]);
 }
@@ -82,7 +82,7 @@
  ******************************************************************************/
 void cm_set_context_by_mpidr(uint64_t mpidr, void *context, uint32_t security_state)
 {
-	assert(security_state <= NON_SECURE);
+	assert(sec_state_is_valid(security_state));
 
 	set_cpu_data_by_mpidr(mpidr, cpu_context[security_state], context);
 }
diff --git a/bl31/interrupt_mgmt.c b/bl31/interrupt_mgmt.c
index 2b0c797..e595634 100644
--- a/bl31/interrupt_mgmt.c
+++ b/bl31/interrupt_mgmt.c
@@ -107,7 +107,7 @@
 {
 	uint32_t scr_el3;
 
-	assert(security_state <= NON_SECURE);
+	assert(sec_state_is_valid(security_state));
 	scr_el3 = intr_type_descs[INTR_TYPE_NS].scr_el3[security_state];
 	scr_el3 |= intr_type_descs[INTR_TYPE_S_EL1].scr_el3[security_state];
 	scr_el3 |= intr_type_descs[INTR_TYPE_EL3].scr_el3[security_state];