From d1a6a71693c31ec1f390671b605183fb58259715 Mon Sep 17 00:00:00 2001
From: Wade Smith <wade.cameron@gmail.com>
Date: Mon, 1 Aug 2022 11:35:00 +1000
Subject: [PATCH] Add special key handling for MEMORY USAGE

Add special handling for the MEMORY USAGE command to extract the key
from the second argument.
---
 action.go      | 2 ++
 action_test.go | 5 +++++
 2 files changed, 7 insertions(+)

diff --git a/action.go b/action.go
index b53ed22..f793de1 100644
--- a/action.go
+++ b/action.go
@@ -233,6 +233,8 @@ func DefaultActionProperties(cmd string, args ...string) ActionProperties {
 	case noKeyCmds[cmd] || len(args) == 0:
 	case cmd == "BITOP" && len(args) > 1:
 		properties.Keys = args[1:]
+	case cmd == "MEMORY" && len(args) > 1 && strings.ToUpper(args[0]) == "USAGE":
+		properties.Keys = args[1:2]
 	case cmd == "MSET":
 		properties.Keys = keysFromKeyValuePairs(args)
 	case cmd == "XINFO":
diff --git a/action_test.go b/action_test.go
index 8921ec7..c054240 100644
--- a/action_test.go
+++ b/action_test.go
@@ -44,6 +44,11 @@ func TestCmdAction(t *T) {
 	var dstval string
 	require.Nil(t, c.Do(ctx, Cmd(&dstval, "GET", key+key)))
 	assert.Equal(t, val, dstval)
+
+	// MEMORY USAGE needs special handling
+	memoryUsageCmd := Cmd(nil, "MEMORY", "USAGE", key)
+	assert.Equal(t, []string{key}, memoryUsageCmd.Properties().Keys)
+	require.Nil(t, c.Do(ctx, memoryUsageCmd))
 }
 
 func TestCmdActionMSet(t *T) {