diff --git a/ci.sh b/ci.sh
index 920c593..a3ebbc2 100755
--- a/ci.sh
+++ b/ci.sh
@@ -31,8 +31,10 @@ fi
 #############
 # Run tests #
 #############
-RUST_BACKTRACE=1 cargo test
+RUST_BACKTRACE=1 cargo test -- --test-threads=1
 
+# Remove mbedtls directory if it exists
+rm -rf psa-crypto/mbedtls
 ################################
 # Check feature configurations #
 ################################
diff --git a/psa-crypto-sys/src/c/shim.c b/psa-crypto-sys/src/c/shim.c
index a5b9846..7829a2c 100644
--- a/psa-crypto-sys/src/c/shim.c
+++ b/psa-crypto-sys/src/c/shim.c
@@ -236,3 +236,21 @@ shim_PSA_KEY_TYPE_DH_PUBLIC_KEY(psa_dh_group_t group)
 {
 	return PSA_KEY_TYPE_DH_PUBLIC_KEY(group);
 }
+
+psa_key_type_t
+shim_PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(psa_key_type_t key_type)
+{
+    return PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(key_type);
+}
+
+size_t
+shim_PSA_SIGN_OUTPUT_SIZE(psa_key_type_t key_type, size_t key_bits, psa_algorithm_t alg)
+{
+    return PSA_SIGN_OUTPUT_SIZE(key_type, key_bits, alg);
+}
+
+size_t
+shim_PSA_KEY_EXPORT_MAX_SIZE(psa_key_type_t key_type, size_t key_bits)
+{
+    return PSA_KEY_EXPORT_MAX_SIZE(key_type, key_bits);
+}
\ No newline at end of file
diff --git a/psa-crypto-sys/src/c/shim.h b/psa-crypto-sys/src/c/shim.h
index 6169869..9e105b4 100644
--- a/psa-crypto-sys/src/c/shim.h
+++ b/psa-crypto-sys/src/c/shim.h
@@ -115,3 +115,6 @@ psa_key_type_t shim_PSA_KEY_TYPE_ECC_KEY_PAIR(psa_ecc_curve_t curve);
 psa_key_type_t shim_PSA_KEY_TYPE_ECC_PUBLIC_KEY(psa_ecc_curve_t curve);
 psa_key_type_t shim_PSA_KEY_TYPE_DH_KEY_PAIR(psa_dh_group_t group);
 psa_key_type_t shim_PSA_KEY_TYPE_DH_PUBLIC_KEY(psa_dh_group_t group);
+psa_key_type_t shim_PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(psa_key_type_t key_type);
+size_t shim_PSA_SIGN_OUTPUT_SIZE(psa_key_type_t key_type, size_t key_bits, psa_algorithm_t alg);
+size_t shim_PSA_KEY_EXPORT_MAX_SIZE(psa_key_type_t key_type, size_t key_bits);
\ No newline at end of file
diff --git a/psa-crypto-sys/src/shim_methods.rs b/psa-crypto-sys/src/shim_methods.rs
index 6e5941d..5db058a 100644
--- a/psa-crypto-sys/src/shim_methods.rs
+++ b/psa-crypto-sys/src/shim_methods.rs
@@ -171,3 +171,19 @@ pub fn PSA_KEY_TYPE_DH_KEY_PAIR(group: psa_dh_group_t) -> psa_key_type_t {
 pub fn PSA_KEY_TYPE_DH_PUBLIC_KEY(group: psa_dh_group_t) -> psa_key_type_t {
     unsafe { psa_crypto_binding::shim_PSA_KEY_TYPE_DH_PUBLIC_KEY(group) }
 }
+
+pub unsafe fn PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(key_type: psa_key_type_t) -> psa_key_type_t {
+    psa_crypto_binding::shim_PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(key_type)
+}
+
+pub unsafe fn PSA_SIGN_OUTPUT_SIZE(
+    key_type: psa_key_type_t,
+    key_bits: usize,
+    alg: psa_algorithm_t,
+) -> usize {
+    psa_crypto_binding::shim_PSA_SIGN_OUTPUT_SIZE(key_type, key_bits, alg)
+}
+
+pub unsafe fn PSA_EXPORT_KEY_OUTPUT_SIZE(key_type: psa_key_type_t, key_bits: usize) -> usize {
+    psa_crypto_binding::shim_PSA_KEY_EXPORT_MAX_SIZE(key_type, key_bits)
+}
diff --git a/psa-crypto/src/operations/asym_signature.rs b/psa-crypto/src/operations/asym_signature.rs
index bc5b83d..4ff384b 100644
--- a/psa-crypto/src/operations/asym_signature.rs
+++ b/psa-crypto/src/operations/asym_signature.rs
@@ -43,12 +43,15 @@ use crate::types::status::{Result, Status};
 /// #     0x94, 0x8E, 0x92, 0x50, 0x35, 0xC2, 0x8C, 0x5C, 0x3C, 0xCA, 0xFE, 0x18, 0xE8, 0x81, 0x37, 0x78,
 /// # ];
 /// psa_crypto::init().unwrap();
-/// let mut signature = vec![0; 256];
 /// let my_key = generate(attributes, None).unwrap();
-/// let size = sign_hash(my_key,
-///                      AsymmetricSignature::RsaPkcs1v15Sign {
+/// let alg = AsymmetricSignature::RsaPkcs1v15Sign {
 ///                          hash_alg: Hash::Sha256.into(),
-///                      },
+///                      };
+/// let buffer_size = attributes.sign_output_size(alg).unwrap();
+/// let mut signature = vec![0; buffer_size];
+///
+/// let size = sign_hash(my_key,
+///                      alg,
 ///                      &HASH,
 ///                      &mut signature).unwrap();
 /// signature.resize(size, 0);
@@ -112,10 +115,11 @@ pub fn sign_hash(
 /// #     0x94, 0x8E, 0x92, 0x50, 0x35, 0xC2, 0x8C, 0x5C, 0x3C, 0xCA, 0xFE, 0x18, 0xE8, 0x81, 0x37, 0x78,
 /// # ];
 /// psa_crypto::init().unwrap();
-/// let mut signature = vec![0; 256];
 /// let alg = AsymmetricSignature::RsaPkcs1v15Sign {
 ///     hash_alg: Hash::Sha256.into(),
 /// };
+/// let buffer_size = attributes.sign_output_size(alg).unwrap();
+/// let mut signature = vec![0; buffer_size];
 /// let my_key = generate(attributes, None).unwrap();
 /// let size = sign_hash(my_key,
 ///                      alg,
diff --git a/psa-crypto/src/operations/key_management.rs b/psa-crypto/src/operations/key_management.rs
index 6ecf4aa..f1c2df0 100644
--- a/psa-crypto/src/operations/key_management.rs
+++ b/psa-crypto/src/operations/key_management.rs
@@ -208,7 +208,8 @@ pub fn import(attributes: Attributes, id: Option<u32>, data: &[u8]) -> Result<Id
 /// #     },
 /// # };
 /// psa_crypto::init().unwrap();
-/// let mut data = vec![0; 256];
+/// let buffer_size = attributes.export_public_key_output_size().unwrap();
+/// let mut data = vec![0; buffer_size];
 /// let my_key = key_management::generate(attributes, None).unwrap();
 /// let size = key_management::export_public(my_key, &mut data).unwrap();
 /// data.resize(size, 0);
diff --git a/psa-crypto/src/types/key.rs b/psa-crypto/src/types/key.rs
index 73741c1..022b99c 100644
--- a/psa-crypto/src/types/key.rs
+++ b/psa-crypto/src/types/key.rs
@@ -6,6 +6,8 @@
 #![allow(deprecated)]
 #[cfg(feature = "with-mbed-crypto")]
 use crate::initialized;
+#[cfg(feature = "with-mbed-crypto")]
+use crate::types::algorithm::AsymmetricSignature;
 use crate::types::algorithm::{Algorithm, Cipher};
 #[cfg(feature = "with-mbed-crypto")]
 use crate::types::status::Status;
@@ -305,6 +307,50 @@ impl Attributes {
         get_attributes_res?;
         Ok(attributes?)
     }
+
+    /// Sufficient size for a buffer to export the key, if supported
+    #[cfg(feature = "with-mbed-crypto")]
+    pub fn export_key_output_size(self) -> Result<usize> {
+        Attributes::export_key_output_size_base(self.key_type, self.bits)
+    }
+
+    /// Sufficient size for a buffer to export the public key, if supported
+    #[cfg(feature = "with-mbed-crypto")]
+    pub fn export_public_key_output_size(self) -> Result<usize> {
+        match self.key_type {
+            Type::RsaKeyPair
+            | Type::RsaPublicKey
+            | Type::EccKeyPair { .. }
+            | Type::EccPublicKey { .. }
+            | Type::DhKeyPair { .. }
+            | Type::DhPublicKey { .. } => {
+                let pub_type = self.key_type.key_type_public_key_of_key_pair()?;
+                Attributes::export_key_output_size_base(pub_type, self.bits)
+            }
+            _ => Err(Error::InvalidArgument),
+        }
+    }
+
+    /// Sufficient size for a buffer to export the given key type, if supported
+    #[cfg(feature = "with-mbed-crypto")]
+    fn export_key_output_size_base(key_type: Type, bits: usize) -> Result<usize> {
+        let size =
+            unsafe { psa_crypto_sys::PSA_EXPORT_KEY_OUTPUT_SIZE(key_type.try_into()?, bits) };
+        if size > 0 {
+            Ok(size)
+        } else {
+            Err(Error::NotSupported)
+        }
+    }
+
+    /// Sufficient buffer size for a signature using the given key, if the key is supported
+    #[cfg(feature = "with-mbed-crypto")]
+    pub fn sign_output_size(self, alg: AsymmetricSignature) -> Result<usize> {
+        self.compatible_with_alg(Algorithm::AsymmetricSignature(alg))?;
+        Ok(unsafe {
+            psa_crypto_sys::PSA_SIGN_OUTPUT_SIZE(self.key_type.try_into()?, self.bits, alg.into())
+        })
+    }
 }
 
 /// The lifetime of a key indicates where it is stored and which application and system actions
@@ -407,6 +453,27 @@ impl Type {
             _ => false,
         }
     }
+
+    /// If key is public or key pair, returns the corresponding public key type.
+    #[cfg(feature = "with-mbed-crypto")]
+    pub fn key_type_public_key_of_key_pair(self) -> Result<Type> {
+        match self {
+            Type::RsaKeyPair
+            | Type::RsaPublicKey
+            | Type::EccKeyPair { .. }
+            | Type::EccPublicKey { .. }
+            | Type::DhKeyPair { .. }
+            | Type::DhPublicKey { .. } => {
+                Ok(
+                    unsafe {
+                        psa_crypto_sys::PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(self.try_into()?)
+                    }
+                    .try_into()?,
+                )
+            }
+            _ => Err(Error::InvalidArgument),
+        }
+    }
 }
 
 /// Enumeration of elliptic curve families supported. They are needed to create an ECC key.