diff --git a/tests/unit/test_aes.cpp b/tests/unit/test_aes.cpp
index b6351ddca..82b50e943 100644
--- a/tests/unit/test_aes.cpp
+++ b/tests/unit/test_aes.cpp
@@ -41,7 +41,7 @@ namespace ot {
  */
 void TestMacBeaconFrame(void)
 {
-    uint8_t key[] = {
+    uint8_t rawKey[] = {
         0xc0, 0xc1, 0xc2, 0xc3, 0xc4, 0xc5, 0xc6, 0xc7, 0xc8, 0xc9, 0xca, 0xcb, 0xcc, 0xcd, 0xce, 0xcf,
     };
 
@@ -57,19 +57,34 @@ void TestMacBeaconFrame(void)
                            0xAC, 0x02, 0x05, 0x00, 0x00, 0x00, 0x55, 0xCF, 0x00, 0x00, 0x51, 0x52,
                            0x53, 0x54, 0x22, 0x3B, 0xC1, 0xEC, 0x84, 0x1A, 0xB5, 0x53};
 
+    uint8_t nonce[] = {
+        0xAC, 0xDE, 0x48, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x05, 0x02,
+    };
+
     otInstance    *instance = testInitInstance();
     Crypto::AesCcm aesCcm;
     uint32_t       headerLength  = sizeof(test) - 8;
     uint32_t       payloadLength = 0;
     uint8_t        tagLength     = 8;
 
-    uint8_t nonce[] = {
-        0xAC, 0xDE, 0x48, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x05, 0x02,
-    };
+#if OPENTHREAD_CONFIG_PLATFORM_KEY_REFERENCES_ENABLE
+    Crypto::Key             key;
+    Crypto::Storage::KeyRef keyRef;
+#endif
 
     VerifyOrQuit(instance != nullptr);
 
-    aesCcm.SetKey(key, sizeof(key));
+#if OPENTHREAD_CONFIG_PLATFORM_KEY_REFERENCES_ENABLE
+    SuccessOrQuit(Crypto::Storage::ImportKey(keyRef, Crypto::Storage::kKeyTypeAes, Crypto::Storage::kKeyAlgorithmAesEcb,
+                                             Crypto::Storage::kUsageEncrypt, Crypto::Storage::kTypeVolatile, rawKey,
+                                             sizeof(rawKey)));
+
+    key.SetAsKeyRef(keyRef);
+    aesCcm.SetKey(key);
+#else
+    aesCcm.SetKey(rawKey, sizeof(rawKey));
+#endif
+
     aesCcm.Init(headerLength, payloadLength, tagLength, nonce, sizeof(nonce));
     aesCcm.Header(test, headerLength);
     VerifyOrQuit(aesCcm.GetTagLength() == tagLength);
@@ -84,6 +99,10 @@ void TestMacBeaconFrame(void)
 
     VerifyOrQuit(memcmp(test, decrypted, sizeof(decrypted)) == 0);
 
+#if OPENTHREAD_CONFIG_PLATFORM_KEY_REFERENCES_ENABLE
+    Crypto::Storage::DestroyKey(keyRef);
+#endif
+
     testFreeInstance(instance);
 }
 
@@ -92,7 +111,7 @@ void TestMacBeaconFrame(void)
  */
 void TestMacCommandFrame(void)
 {
-    uint8_t key[] = {
+    uint8_t rawKey[] = {
         0xc0, 0xc1, 0xc2, 0xc3, 0xc4, 0xc5, 0xc6, 0xc7, 0xc8, 0xc9, 0xca, 0xcb, 0xcc, 0xcd, 0xce, 0xcf,
     };
 
@@ -128,9 +147,24 @@ void TestMacCommandFrame(void)
     Message       *message;
     Crypto::AesCcm aesCcm;
 
+#if OPENTHREAD_CONFIG_PLATFORM_KEY_REFERENCES_ENABLE
+    Crypto::Key             key;
+    Crypto::Storage::KeyRef keyRef;
+#endif
+
     VerifyOrQuit(instance != nullptr);
 
-    aesCcm.SetKey(key, sizeof(key));
+#if OPENTHREAD_CONFIG_PLATFORM_KEY_REFERENCES_ENABLE
+    SuccessOrQuit(Crypto::Storage::ImportKey(keyRef, Crypto::Storage::kKeyTypeAes, Crypto::Storage::kKeyAlgorithmAesEcb,
+                                             Crypto::Storage::kUsageEncrypt, Crypto::Storage::kTypeVolatile, rawKey,
+                                             sizeof(rawKey)));
+
+    key.SetAsKeyRef(keyRef);
+    aesCcm.SetKey(key);
+#else
+    aesCcm.SetKey(rawKey, sizeof(rawKey));
+#endif
+
     aesCcm.Init(kHeaderLength, kPayloadLength, kTagLength, nonce, sizeof(nonce));
     aesCcm.Header(test, kHeaderLength);
     aesCcm.Payload(test + kHeaderLength, test + kHeaderLength, kPayloadLength, Crypto::AesCcm::kEncrypt);
@@ -171,6 +205,11 @@ void TestMacCommandFrame(void)
     VerifyOrQuit(message->Compare(0, decrypted));
 
     message->Free();
+
+#if OPENTHREAD_CONFIG_PLATFORM_KEY_REFERENCES_ENABLE
+    Crypto::Storage::DestroyKey(keyRef);
+#endif
+
     testFreeInstance(instance);
 }
 
@@ -182,7 +221,7 @@ void TestInPlaceAesCcmProcessing(void)
     static constexpr uint16_t kTagLength    = 4;
     static constexpr uint32_t kHeaderLength = 19;
 
-    static const uint8_t kKey[] = {
+    static const uint8_t kRawKey[] = {
         0xa0, 0xa1, 0xa2, 0xa3, 0xa4, 0xa5, 0xa6, 0xa7, 0xa8, 0xa9, 0xaa, 0xab, 0xac, 0xad, 0xae, 0xaf,
     };
 
@@ -200,12 +239,26 @@ void TestInPlaceAesCcmProcessing(void)
     Message       *message;
     Message       *messageClone;
 
+#if OPENTHREAD_CONFIG_PLATFORM_KEY_REFERENCES_ENABLE
+    Crypto::Key             key;
+    Crypto::Storage::KeyRef keyRef;
+#endif
+
     VerifyOrQuit(instance != nullptr);
 
     message = instance->Get<MessagePool>().Allocate(Message::kTypeIp6);
     VerifyOrQuit(message != nullptr);
 
-    aesCcm.SetKey(kKey, sizeof(kKey));
+#if OPENTHREAD_CONFIG_PLATFORM_KEY_REFERENCES_ENABLE
+    SuccessOrQuit(Crypto::Storage::ImportKey(keyRef, Crypto::Storage::kKeyTypeAes, Crypto::Storage::kKeyAlgorithmAesEcb,
+                                             Crypto::Storage::kUsageEncrypt, Crypto::Storage::kTypeVolatile, kRawKey,
+                                             sizeof(kRawKey)));
+
+    key.SetAsKeyRef(keyRef);
+    aesCcm.SetKey(key);
+#else
+    aesCcm.SetKey(kRawKey, sizeof(kRawKey));
+#endif
 
     for (uint16_t msgLength : kMessageLengths)
     {
@@ -251,6 +304,11 @@ void TestInPlaceAesCcmProcessing(void)
     }
 
     message->Free();
+
+#if OPENTHREAD_CONFIG_PLATFORM_KEY_REFERENCES_ENABLE
+    Crypto::Storage::DestroyKey(keyRef);
+#endif
+
     testFreeInstance(instance);
 }