Fix add-keymap-based-bindings and associated test

2 files changed, 18 insertions, 19 deletions

diff --git a/which-key-tests.el b/which-key-tests.el
index eeedb557370..04617377999 100644
--- a/which-key-tests.el
+++ b/which-key-tests.el
@@ -29,20 +29,17 @@
 
 (ert-deftest which-key-test--keymap-based-bindings ()
   (let ((map (make-sparse-keymap))
-        (emacs-lisp-mode-map (copy-keymap emacs-lisp-mode-map)))
-    (emacs-lisp-mode)
-    (define-key map "x" 'ignore)
-    (define-key emacs-lisp-mode-map "\C-c\C-a" 'complete)
-    (define-key emacs-lisp-mode-map "\C-c\C-b" map)
-    (which-key-add-keymap-based-replacements emacs-lisp-mode-map
-      "C-c C-a" '("mycomplete" . complete)
-      "C-c C-b" "mymap")
-    (should (equal
-             (which-key--maybe-replace '("C-c C-a" . "complete"))
-             '("C-c C-a" . "mycomplete")))
-    (should (equal
-             (which-key--maybe-replace '("C-c C-b" . ""))
-             '("C-c C-b" . "mymap")))))
+        (prefix-map (make-sparse-keymap)))
+    (define-key prefix-map "x" 'ignore)
+    (define-key map "\C-a" 'complete)
+    (define-key map "\C-b" prefix-map)
+    (which-key-add-keymap-based-replacements map
+      "C-a" '("mycomplete" . complete)
+      "C-b" "mymap")
+    (should (equal
+             (which-key--get-keymap-bindings map)
+             '(("C-a" . "mycomplete")
+               ("C-b" . "mymap"))))))
 
 (ert-deftest which-key-test--prefix-declaration ()
   "Test `which-key-declare-prefixes' and
diff --git a/which-key.el b/which-key.el
index d6baa70b537..ec3f760159b 100644
--- a/which-key.el
+++ b/which-key.el
@@ -914,11 +914,13 @@ both have the same effect for the \"C-x C-w\" key binding, but
 the latter causes which-key to verify that the key sequence is
 actually bound to write-file before performing the replacement."
   (while key
-    (let ((string (if (stringp replacement)
-                      replacement
-                    (car-safe replacement)))
-          (command (cdr-safe replacement)))
-      (define-key keymap (kbd key) (cons string command)))
+    (cond ((consp replacement)
+           (define-key keymap (kbd key) replacement))
+          ((stringp replacement)
+           (define-key keymap (kbd key) (cons replacement
+                                              (lookup-key keymap (kbd key)))))
+          (t
+           (user-error "replacement is neither a cons cell or a string")))
     (setq key (pop more)
           replacement (pop more))))
 (put 'which-key-add-keymap-based-replacements 'lisp-indent-function 'defun)