diff --git a/src/genkey/clu_genkey.c b/src/genkey/clu_genkey.c
index 178f0f9..6ed18f2 100644
--- a/src/genkey/clu_genkey.c
+++ b/src/genkey/clu_genkey.c
@@ -1039,7 +1039,7 @@ int wolfCLU_genKey_RSA(WC_RNG* rng, char* fName, int directive, int fmt, int
 int wolfCLU_genKey_Dilithium(WC_RNG* rng, char* fName, int directive, int fmt,
                             int keySz, int level, int withAlg)
 {
-#ifdef HAVE_DILITHIUM    
+#ifdef HAVE_DILITHIUM
     int    ret = WOLFCLU_SUCCESS;
 
     XFILE  file = NULL;
@@ -1076,11 +1076,11 @@ int wolfCLU_genKey_Dilithium(WC_RNG* rng, char* fName, int directive, int fmt,
 
     /* init the dilithium key */
     if (wc_dilithium_init(key) != 0) {
-        wolfCLU_LogError("Failed to initialize Dilithium Key.\nRET: %d", ret);
+        wolfCLU_LogError("Failed to initialize Dilithium Key.");
     #ifdef WOLFSSL_SMALL_STACK
         XFREE(key, HEAP_HINT, DYNAMIC_TYPE_DILITHIUM);
     #endif
-        return ret;
+        return WOLFCLU_FATAL_ERROR;
     }
 
     /* set the level of the dilithium key */
@@ -1113,12 +1113,6 @@ int wolfCLU_genKey_Dilithium(WC_RNG* rng, char* fName, int directive, int fmt,
     if (ret == WOLFCLU_SUCCESS) {
         XMEMSET(fOutNameBuf, 0, fNameSz + fExtSz);
         XMEMCPY(fOutNameBuf, fName, fNameSz);
-
-        derBuf = (byte*)XMALLOC(keySz, HEAP_HINT,
-                DYNAMIC_TYPE_TMP_BUFFER);
-        if (derBuf == NULL) {
-            ret = MEMORY_E;
-        }
     }
 
     if (ret == WOLFCLU_SUCCESS) {
@@ -1129,6 +1123,13 @@ int wolfCLU_genKey_Dilithium(WC_RNG* rng, char* fName, int directive, int fmt,
             case PRIV_ONLY_FILE:
                 /* add on the final part of the file name ".priv" */
                 XMEMCPY(fOutNameBuf + fNameSz, fExtPriv, fExtSz);
+
+                derBuf = (byte*)XMALLOC(keySz, HEAP_HINT,
+                        DYNAMIC_TYPE_TMP_BUFFER);
+                if (derBuf == NULL) {
+                    ret = MEMORY_E;
+                    break;
+                }
                 WOLFCLU_LOG(WOLFCLU_L0, "Private key file = %s", fOutNameBuf);
 
                 /* Private key to der */
@@ -1184,7 +1185,7 @@ int wolfCLU_genKey_Dilithium(WC_RNG* rng, char* fName, int directive, int fmt,
 
                 FALL_THROUGH;
             case PUB_ONLY_FILE:
-                /* add on the final part of the file name ".priv" */
+                /* add on the final part of the file name ".pub" */
                 XMEMCPY(fOutNameBuf + fNameSz, fExtPub, fExtSz);
                 WOLFCLU_LOG(WOLFCLU_L0, "Public key file = %s", fOutNameBuf);
 
@@ -1192,8 +1193,9 @@ int wolfCLU_genKey_Dilithium(WC_RNG* rng, char* fName, int directive, int fmt,
                                         DYNAMIC_TYPE_TMP_BUFFER);
                 if (derBuf == NULL) {
                     ret = MEMORY_E;
+                    break;
                 }
-                
+
                 derBufSz = wc_Dilithium_PublicKeyToDer(key, derBuf,
                                                 (word32)keySz, withAlg);
                 if (derBufSz < 0) {
diff --git a/tests/genkey_sign_ver/genkey-sign-ver-test.sh b/tests/genkey_sign_ver/genkey-sign-ver-test.sh
index ee68739..7334a84 100755
--- a/tests/genkey_sign_ver/genkey-sign-ver-test.sh
+++ b/tests/genkey_sign_ver/genkey-sign-ver-test.sh
@@ -40,6 +40,10 @@ cleanup_genkey_sign_ver(){
     rm rsakey.pub
     rm mldsakey.priv
     rm mldsakey.pub
+    rm mldsakey_pub.pub
+    rm mldsakey_pub.priv
+    rm mldsakey_priv.pub
+    rm mldsakey_priv.priv
     rm ecc-signed.sig
     rm ed-signed.sig
     rm rsa-signed.sig
@@ -89,7 +93,7 @@ rsa_compare_decrypted(){
     else
         printf '%s\n' "Decrypted mismatch with original, FAILURE!"
         printf '%s\n' "DECRYPTED --> ${1}"
-        printf '%s\n' "ORIGINAL --> ${2}" && exit -1
+        printf '%s\n' "ORIGINAL --> ${2}" && exit 99
     fi
 }
 
@@ -109,13 +113,13 @@ gen_key_sign_ver_test(){
     printf '%s\n' "genkey RESULT - $RESULT"
     [ $RESULT -ne 0 ] && printf '%s\n' "Failed $1 genkey" && \
     printf '%s\n' "Before running this test please configure wolfssl with" && \
-    printf '%s\n' "--enable-keygen" && exit -1
+    printf '%s\n' "--enable-keygen" && exit 99
 
     # test signing with priv key
     ./wolfssl -$1 -sign -inkey $2.priv -inform $4 -in sign-this.txt -out $3
     RESULT=$?
     printf '%s\n' "sign RESULT - $RESULT"
-    [ $RESULT -ne 0 ] && printf '%s\n' "Failed $1 sign" && exit -1
+    [ $RESULT -ne 0 ] && printf '%s\n' "Failed $1 sign" && exit 99
 
     # test verifying with priv key
     if [ "${1}" = "rsa" ]; then
@@ -130,7 +134,7 @@ gen_key_sign_ver_test(){
     fi
     RESULT=$?
     printf '%s\n' "private verify RESULT - $RESULT"
-    [ $RESULT -ne 0 ] && printf '%s\n' "Failed $1 private verify" && exit -1
+    [ $RESULT -ne 0 ] && printf '%s\n' "Failed $1 private verify" && exit 99
 
     # test verifying with pub key
     if [ "${1}" = "rsa" ]; then
@@ -141,7 +145,7 @@ gen_key_sign_ver_test(){
     fi
     RESULT=$?
     printf '%s\n' "public verify RESULT - $RESULT"
-    [ $RESULT -ne 0 ] && printf '%s\n' "Failed $1 public verify " && exit -1
+    [ $RESULT -ne 0 ] && printf '%s\n' "Failed $1 public verify " && exit 99
 
     if [ $1 = "rsa" ]; then
         ORIGINAL=`cat -A sign-this.txt`
@@ -227,6 +231,20 @@ for level in 2 3 5
 do
     gen_key_sign_ver_test ${ALGORITHM} ${KEYFILENAME} ${SIGOUTNAME} ${DERPEMRAW} ${level}
 done
+
+# Verifies that -output PUB generates only the public key file.
+./wolfssl -genkey dilithium -level 2 -out mldsakey_pub -outform der -output pub
+RESULT=$?
+[ $RESULT -ne 0 ] && printf '%s\n' "Failed dilithium genkey -output PUB" && exit 99
+[ ! -f mldsakey_pub.pub ]  && printf '%s\n' "dilithium -output PUB: .pub file missing"           && exit 99
+[ -f  mldsakey_pub.priv ]  && printf '%s\n' "dilithium -output PUB: .priv unexpectedly created"  && exit 99
+
+# Verifies that -output PRIV generates only the private key file.
+./wolfssl -genkey dilithium -level 2 -out mldsakey_priv -outform der -output priv
+RESULT=$?
+[ $RESULT -ne 0 ] && printf '%s\n' "Failed dilithium genkey -output PRIV" && exit 99
+[ ! -f mldsakey_priv.priv ] && printf '%s\n' "dilithium -output PRIV: .priv file missing"         && exit 99
+[ -f  mldsakey_priv.pub ]   && printf '%s\n' "dilithium -output PRIV: .pub unexpectedly created"  && exit 99
 fi
 
 # Check if xmss is availabe