diff --git a/tests/suites/test_suite_aes.function b/tests/suites/test_suite_aes.function
index 91f5fa2..e998795 100644
--- a/tests/suites/test_suite_aes.function
+++ b/tests/suites/test_suite_aes.function
@@ -161,20 +161,18 @@
     unsigned char src_str[100] = { 0, };
     unsigned char dst_str[100] = { 0, };
     unsigned char output[100]  = { 0, };
-    mbedtls_aes_context crypt_ctx, tweak_ctx;
+    mbedtls_aes_xts_context ctx;
     int key_len, data_len;
 
-    mbedtls_aes_init( &crypt_ctx );
-    mbedtls_aes_init( &tweak_ctx );
+    mbedtls_aes_xts_init( &ctx );
 
     key_len = unhexify( key_str, hex_key_string );
     unhexify( iv_str, hex_iv_string );
     data_len = unhexify( src_str, hex_src_string );
 
-    mbedtls_aes_setkey_enc( &crypt_ctx, key_str,               ( key_len * 8 ) / 2 );
-    mbedtls_aes_setkey_enc( &tweak_ctx, key_str + key_len / 2, ( key_len * 8 ) / 2 );
+    mbedtls_aes_xts_setkey_enc( &ctx, key_str, key_len * 8 );
 
-    TEST_ASSERT( mbedtls_aes_crypt_xts( &crypt_ctx, &tweak_ctx, MBEDTLS_AES_ENCRYPT, data_unit_len, iv_str, src_str, output ) == xts_result );
+    TEST_ASSERT( mbedtls_aes_crypt_xts( &ctx, MBEDTLS_AES_ENCRYPT, data_unit_len, iv_str, src_str, output ) == xts_result );
     if( xts_result == 0 )
     {
         hexify( dst_str, output, data_len );
@@ -183,8 +181,7 @@
     }
 
 exit:
-    mbedtls_aes_free( &crypt_ctx );
-    mbedtls_aes_free( &tweak_ctx );
+    mbedtls_aes_xts_free( &ctx );
 }
 /* END_CASE */
 
@@ -198,20 +195,18 @@
     unsigned char src_str[100] = { 0, };
     unsigned char dst_str[100] = { 0, };
     unsigned char output[100]  = { 0, };
-    mbedtls_aes_context crypt_ctx, tweak_ctx;
+    mbedtls_aes_xts_context ctx;
     int key_len, data_len;
 
-    mbedtls_aes_init( &crypt_ctx );
-    mbedtls_aes_init( &tweak_ctx );
+    mbedtls_aes_xts_init( &ctx );
 
     key_len = unhexify( key_str, hex_key_string );
     unhexify( iv_str, hex_iv_string );
     data_len = unhexify( src_str, hex_src_string );
 
-    mbedtls_aes_setkey_dec( &crypt_ctx, key_str,               ( key_len * 8 ) / 2 );
-    mbedtls_aes_setkey_enc( &tweak_ctx, key_str + key_len / 2, ( key_len * 8 ) / 2 );
+    mbedtls_aes_xts_setkey_dec( &ctx, key_str, key_len * 8 );
 
-	TEST_ASSERT( mbedtls_aes_crypt_xts( &crypt_ctx, &tweak_ctx, MBEDTLS_AES_DECRYPT, data_unit_len, iv_str, src_str, output ) == xts_result );
+    TEST_ASSERT( mbedtls_aes_crypt_xts( &ctx, MBEDTLS_AES_DECRYPT, data_unit_len, iv_str, src_str, output ) == xts_result );
     if( xts_result == 0 )
     {
         hexify( dst_str, output, data_len );
@@ -220,8 +215,7 @@
     }
 
 exit:
-    mbedtls_aes_free( &crypt_ctx );
-    mbedtls_aes_free( &tweak_ctx );
+    mbedtls_aes_xts_free( &ctx );
 }
 /* END_CASE */