diff --git a/lib/kestrel/client.rb b/lib/kestrel/client.rb index 3a2c54e..72038e3 100644 --- a/lib/kestrel/client.rb +++ b/lib/kestrel/client.rb @@ -147,7 +147,9 @@ def extract_kestrel_options!(opts) end def shuffle_if_necessary!(key) - if key != @current_queue || @counter >= @gets_per_server + # Don't reset servers on the first request: + # i.e. @counter == 0 && @current_queue == nil + if (@counter > 0 && key != @current_queue) || @counter >= @gets_per_server @counter = 0 @current_queue = key @read_client.quit diff --git a/spec/kestrel/client_spec.rb b/spec/kestrel/client_spec.rb index e9d3073..49f5cb5 100644 --- a/spec/kestrel/client_spec.rb +++ b/spec/kestrel/client_spec.rb @@ -18,17 +18,21 @@ end it "gets from the same server :gets_per_server times" do - mock(@kestrel).get_from_last("a_queue/t=10", false).times(100) { 'item' } - mock(@kestrel).get_from_random("a_queue/t=10", false).times(2) { 'item' } + client = @kestrel.instance_variable_get(:@read_client) + mock(client).get_from_last("a_queue/t=10", true).times(102).returns('item') + mock(client).quit.once + mock(client).set_servers.with(anything).once 102.times { @kestrel.get("a_queue") } end it "gets from a different server when the last result was nil" do - mock(@kestrel).get_from_last("a_queue/t=10", false).never { nil } - mock(@kestrel).get_from_random("a_queue/t=10", false).times(3) { nil } + client = @kestrel.instance_variable_get(:@read_client) + mock(client).get_from_last("a_queue/t=10", true).returns(nil).twice + mock(client).quit.once + mock(client).set_servers.with(anything).once - 3.times { @kestrel.get("a_queue") } + 2.times { @kestrel.get("a_queue") } end it "returns nil if there is a recoverable exception" do