diff --git a/default.nix b/default.nix
index c4f079c..d87ace3 100644
--- a/default.nix
+++ b/default.nix
@@ -11,11 +11,10 @@ let
   });
 in rec {
   toStorePath = target:
-    # If a store path has been given, transform it to a string. And make sure
-    # to capture the store path in its context.
-    if lib.isStorePath target then
-      let path = toString target; in
-      builtins.appendContext path { "${path}" = { path = true; }; }
+    # If a store path has been given but is a string, add the missing context
+    # to it so it will be propagated properly as a build input.
+    if builtins.isString target && lib.isStorePath target then
+      builtins.appendContext target { "${target}" = { path = true; }; }
     # Otherwise, add to the store. This takes care of appending the store path
     # in the context automatically.
     else "${target}";
diff --git a/test.sh b/test.sh
new file mode 100755
index 0000000..f5a5130
--- /dev/null
+++ b/test.sh
@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+set -euo pipefail
+export NIX_PATH=channel:nixos-20.09
+
+echo "Test with attribute name"
+./nix-bundle.sh hello /bin/hello
+
+echo "Test with store path"
+out=$(nix-build --no-out-link --expr '(import <nixpkgs> {})' -A hello)
+./nix-bundle.sh "$out" /bin/hello