diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index ca461c20e..6dff2f532 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -26,6 +26,20 @@
 #endif
 
 #if defined(MBEDTLS_PSA_CRYPTO_C)
+/*
+ * In case MBEDTLS_PSA_CRYPTO_SPM is defined the code is built for SPM (Secure
+ * Partition Manager) integration which separate the code into two parts
+ * NSPE (Non-Secure Process Environment) and SPE (Secure Process Environment).
+ * In this mode an additional header file should be included.
+ */
+#if defined(MBEDTLS_PSA_CRYPTO_SPM)
+/*
+ * PSA_CRYPTO_SECURE means that this file is compiled to the SPE side.
+ * some headers will be affected by this flag.
+ */
+#define PSA_CRYPTO_SECURE 1
+#include "crypto_spe.h"
+#endif
 
 #include "psa/crypto.h"
 
@@ -2482,6 +2496,7 @@ psa_status_t psa_cipher_abort( psa_cipher_operation_t *operation )
 /* Key Policy */
 /****************************************************************/
 
+#if !defined(MBEDTLS_PSA_CRYPTO_SPM)
 void psa_key_policy_init( psa_key_policy_t *policy )
 {
     memset( policy, 0, sizeof( *policy ) );
@@ -2504,6 +2519,7 @@ psa_algorithm_t psa_key_policy_get_algorithm( psa_key_policy_t *policy )
 {
     return( policy->alg );
 }
+#endif /* !defined(MBEDTLS_PSA_CRYPTO_SPM) */
 
 psa_status_t psa_set_key_policy( psa_key_slot_t key,
                                  const psa_key_policy_t *policy )
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 1017e88c8..c90447f81 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -1,5 +1,10 @@
 /* BEGIN_HEADER */
 #include <stdint.h>
+
+#if defined(MBEDTLS_PSA_CRYPTO_SPM)
+#include "spm/psa_defs.h"
+#endif
+
 #include "mbedtls/asn1write.h"
 #include "psa/crypto.h"
 
@@ -1057,7 +1062,8 @@ void cipher_encrypt( int alg_arg, int key_type_arg,
 
     TEST_ASSERT( psa_encrypt_set_iv( &operation,
                                      iv, iv_size ) == PSA_SUCCESS );
-    output_buffer_size = input->len + operation.block_size;
+    output_buffer_size = (size_t) input->len +
+                         PSA_BLOCK_CIPHER_BLOCK_SIZE( key_type );
     output = mbedtls_calloc( 1, output_buffer_size );
     TEST_ASSERT( output != NULL );
 
@@ -1126,7 +1132,8 @@ void cipher_encrypt_multipart( int alg_arg, int key_type_arg,
 
     TEST_ASSERT( psa_encrypt_set_iv( &operation,
                                      iv, sizeof( iv ) ) == PSA_SUCCESS );
-    output_buffer_size = input->len + operation.block_size;
+    output_buffer_size = (size_t) input->len +
+                         PSA_BLOCK_CIPHER_BLOCK_SIZE( key_type );
     output = mbedtls_calloc( 1, output_buffer_size );
     TEST_ASSERT( output != NULL );
 
@@ -1199,7 +1206,8 @@ void cipher_decrypt_multipart( int alg_arg, int key_type_arg,
     TEST_ASSERT( psa_encrypt_set_iv( &operation,
                                      iv, sizeof( iv ) ) == PSA_SUCCESS );
 
-    output_buffer_size = input->len + operation.block_size;
+    output_buffer_size = (size_t) input->len +
+                         PSA_BLOCK_CIPHER_BLOCK_SIZE( key_type );
     output = mbedtls_calloc( 1, output_buffer_size );
     TEST_ASSERT( output != NULL );
 
@@ -1273,7 +1281,8 @@ void cipher_decrypt( int alg_arg, int key_type_arg,
     TEST_ASSERT( psa_encrypt_set_iv( &operation,
                                      iv, iv_size ) == PSA_SUCCESS );
 
-    output_buffer_size = input->len + operation.block_size;
+    output_buffer_size = (size_t) input->len +
+                         PSA_BLOCK_CIPHER_BLOCK_SIZE( key_type );
     output = mbedtls_calloc( 1, output_buffer_size );
     TEST_ASSERT( output != NULL );
 
@@ -1343,7 +1352,8 @@ void cipher_verify_output( int alg_arg, int key_type_arg,
     TEST_ASSERT( psa_encrypt_generate_iv( &operation1,
                                           iv, iv_size,
                                           &iv_length ) == PSA_SUCCESS );
-    output1_size = input->len + operation1.block_size;
+    output1_size = (size_t) input->len +
+                   PSA_BLOCK_CIPHER_BLOCK_SIZE( key_type );
     output1 = mbedtls_calloc( 1, output1_size );
     TEST_ASSERT( output1 != NULL );
 
@@ -1429,7 +1439,8 @@ void cipher_verify_output_multipart( int alg_arg,
     TEST_ASSERT( psa_encrypt_generate_iv( &operation1,
                                           iv, iv_size,
                                           &iv_length ) == PSA_SUCCESS );
-    output1_buffer_size = input->len + operation1.block_size;
+    output1_buffer_size = (size_t) input->len +
+                          PSA_BLOCK_CIPHER_BLOCK_SIZE( key_type );
     output1 = mbedtls_calloc( 1, output1_buffer_size );
     TEST_ASSERT( output1 != NULL );