@@ -20,6 +20,7 @@ type shellServerState struct {
2020 sawAuth bool
2121 commands []string
2222 blocks map [string ]string
23+ present map [string ]bool
2324}
2425
2526func (s * shellServerState ) record (cmd string ) {
@@ -37,7 +38,17 @@ func (s *shellServerState) setAuth() {
3738func (s * shellServerState ) setBlock (x string , y string , bits string ) {
3839 s .mu .Lock ()
3940 defer s .mu .Unlock ()
40- s .blocks [x + ":" + y ] = bits
41+ key := x + ":" + y
42+ s .blocks [key ] = bits
43+ s .present [key ] = true
44+ }
45+
46+ func (s * shellServerState ) unsetBlock (x string , y string ) {
47+ s .mu .Lock ()
48+ defer s .mu .Unlock ()
49+ key := x + ":" + y
50+ delete (s .present , key )
51+ s .blocks [key ] = "0000"
4152}
4253
4354func (s * shellServerState ) getBlock (x string , y string ) string {
@@ -49,6 +60,12 @@ func (s *shellServerState) getBlock(x string, y string) string {
4960 return "0000"
5061}
5162
63+ func (s * shellServerState ) blockExists (x string , y string ) bool {
64+ s .mu .Lock ()
65+ defer s .mu .Unlock ()
66+ return s .present [x + ":" + y ]
67+ }
68+
5269func startShellTestServer (t * testing.T , token string ) (string , * shellServerState , func ()) {
5370 t .Helper ()
5471
@@ -57,7 +74,10 @@ func startShellTestServer(t *testing.T, token string) (string, *shellServerState
5774 t .Fatalf ("listen: %v" , err )
5875 }
5976
60- state := & shellServerState {blocks : make (map [string ]string )}
77+ state := & shellServerState {
78+ blocks : make (map [string ]string ),
79+ present : make (map [string ]bool ),
80+ }
6181 done := make (chan struct {})
6282
6383 go func () {
@@ -138,6 +158,23 @@ func startShellTestServer(t *testing.T, token string) (string, *shellServerState
138158 if err := writeSimple (writer , "OK" ); err != nil {
139159 return
140160 }
161+ case "UNSET" :
162+ if ! authed {
163+ if err := writeError (writer , "AUTH_REQUIRED use AUTH <token>" ); err != nil {
164+ return
165+ }
166+ continue
167+ }
168+ if len (fields ) != 3 {
169+ if err := writeError (writer , "INVALID_ARGUMENT UNSET requires 2 args" ); err != nil {
170+ return
171+ }
172+ continue
173+ }
174+ state .unsetBlock (fields [1 ], fields [2 ])
175+ if err := writeSimple (writer , "OK" ); err != nil {
176+ return
177+ }
141178 case "GET" :
142179 if ! authed {
143180 if err := writeError (writer , "AUTH_REQUIRED use AUTH <token>" ); err != nil {
@@ -154,6 +191,28 @@ func startShellTestServer(t *testing.T, token string) (string, *shellServerState
154191 if err := writeBulk (writer , []byte (state .getBlock (fields [1 ], fields [2 ]))); err != nil {
155192 return
156193 }
194+ case "EXISTS" :
195+ if ! authed {
196+ if err := writeError (writer , "AUTH_REQUIRED use AUTH <token>" ); err != nil {
197+ return
198+ }
199+ continue
200+ }
201+ if len (fields ) != 3 {
202+ if err := writeError (writer , "INVALID_ARGUMENT EXISTS requires 2 args" ); err != nil {
203+ return
204+ }
205+ continue
206+ }
207+ if state .blockExists (fields [1 ], fields [2 ]) {
208+ if err := writeSimple (writer , "1" ); err != nil {
209+ return
210+ }
211+ } else {
212+ if err := writeSimple (writer , "0" ); err != nil {
213+ return
214+ }
215+ }
157216 case "INFO" :
158217 if ! authed {
159218 if err := writeError (writer , "AUTH_REQUIRED use AUTH <token>" ); err != nil {
@@ -241,7 +300,7 @@ func isBits(bits string) bool {
241300 return bits != ""
242301}
243302
244- func TestRunShellConnectAuthPingGetSetQuit (t * testing.T ) {
303+ func TestRunShellConnectAuthPingExistsGetSetUnsetQuit (t * testing.T ) {
245304 uri , state , stop := startShellTestServer (t , "dev-token" )
246305 defer stop ()
247306
@@ -256,7 +315,7 @@ func TestRunShellConnectAuthPingGetSetQuit(t *testing.T) {
256315 }
257316 defer client .Close ()
258317
259- input := strings .NewReader ("ping\n set 1 2 1010\n get 1 2\n quit\n " )
318+ input := strings .NewReader ("ping\n exists 1 2 \ n\n exists 1 2 \n get 1 2 \n unset 1 2 \n exists 1 2 \n get 1 2\n quit\n " )
260319 var out bytes.Buffer
261320 var errOut bytes.Buffer
262321
@@ -270,16 +329,26 @@ func TestRunShellConnectAuthPingGetSetQuit(t *testing.T) {
270329
271330 state .mu .Lock ()
272331 sawAuth := state .sawAuth
332+ commands := append ([]string (nil ), state .commands ... )
273333 state .mu .Unlock ()
274334 if ! sawAuth {
275335 t .Fatalf ("expected shell auto-auth to run AUTH" )
276336 }
337+ expectedCommands := []string {"AUTH" , "PING" , "EXISTS" , "SET" , "EXISTS" , "GET" , "UNSET" , "EXISTS" , "GET" , "QUIT" }
338+ if len (commands ) != len (expectedCommands ) {
339+ t .Fatalf ("unexpected command count: got %v want %v" , commands , expectedCommands )
340+ }
341+ for i := range expectedCommands {
342+ if commands [i ] != expectedCommands [i ] {
343+ t .Fatalf ("unexpected command sequence: got %v want %v" , commands , expectedCommands )
344+ }
345+ }
277346
278347 output := out .String ()
279- if strings .Count (output , "chunk> " ) < 4 {
348+ if strings .Count (output , "chunk> " ) < 8 {
280349 t .Fatalf ("expected repeated prompt, got %q" , output )
281350 }
282- for _ , expected := range []string {"PONG" , "OK " , "1010" , "BYE" } {
351+ for _ , expected := range []string {"PONG" , "chunk> 0 \n " , "chunk> 1 \n " , " 1010" , "0000 ""BYE" } {
283352 if ! strings .Contains (output , expected ) {
284353 t .Fatalf ("expected %q in output %q" , expected , output )
285354 }
0 commit comments