File: redis_mock.rb

package info (click to toggle)
ruby-redis 4.2.5-1
  • links: PTS, VCS
  • area: main
  • in suites: bullseye
  • size: 1,168 kB
  • sloc: ruby: 12,820; makefile: 107; sh: 24
file content (130 lines) | stat: -rw-r--r-- 3,156 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
# frozen_string_literal: true

require "socket"

module RedisMock
  class Server
    def initialize(options = {})
      tcp_server = TCPServer.new(options[:host] || "127.0.0.1", 0)
      tcp_server.setsockopt(Socket::SOL_SOCKET, Socket::SO_REUSEADDR, true)

      if options[:ssl]
        ctx = OpenSSL::SSL::SSLContext.new

        ssl_params = options.fetch(:ssl_params, {})
        ctx.set_params(ssl_params) unless ssl_params.empty?

        @server = OpenSSL::SSL::SSLServer.new(tcp_server, ctx)
      else
        @server = tcp_server
      end
    end

    def port
      @server.addr[1]
    end

    def start(&block)
      @thread = Thread.new { run(&block) }
    end

    def shutdown
      @thread.kill
    end

    def run
      loop do
        session = @server.accept

        begin
          return if yield(session) == :exit
        ensure
          session.close
        end
      end
    rescue => ex
      warn "Error running mock server: #{ex.class}: #{ex.message}"
      warn ex.backtrace
      retry
    ensure
      @server.close
    end
  end

  # Starts a mock Redis server in a thread.
  #
  # The server will use the lambda handler passed as argument to handle
  # connections. For example:
  #
  #   handler = lambda { |session| session.close }
  #   RedisMock.start_with_handler(handler) do
  #     # Every connection will be closed immediately
  #   end
  #
  def self.start_with_handler(blk, options = {})
    server = Server.new(options)
    port = server.port

    begin
      server.start(&blk)
      yield(port)
    ensure
      server.shutdown
    end
  end

  # Starts a mock Redis server in a thread.
  #
  # The server will reply with a `+OK` to all commands, but you can
  # customize it by providing a hash. For example:
  #
  #   RedisMock.start(:ping => lambda { "+PONG" }) do |port|
  #     assert_equal "PONG", Redis.new(:port => port).ping
  #   end
  #
  def self.start(commands, options = {}, &blk)
    handler = lambda do |session|
      while line = session.gets
        argv = Array.new(line[1..-3].to_i) do
          bytes = session.gets[1..-3].to_i
          arg = session.read(bytes)
          session.read(2) # Discard \r\n
          arg
        end

        command = argv.shift
        blk = commands[command.to_sym]
        blk ||= ->(*_) { "+OK" }

        response = blk.call(*argv)

        # Convert a nil response to :close
        response ||= :close

        if response == :exit
          break :exit
        elsif response == :close
          break :close
        elsif response.is_a?(Array)
          session.write("*%d\r\n" % response.size)

          response.each do |resp|
            if resp.is_a?(Array)
              session.write("*%d\r\n" % resp.size)
              resp.each do |r|
                session.write("$%d\r\n%s\r\n" % [r.length, r])
              end
            else
              session.write("$%d\r\n%s\r\n" % [resp.length, resp])
            end
          end
        else
          session.write(response)
          session.write("\r\n") unless response.end_with?("\r\n")
        end
      end
    end

    start_with_handler(handler, options, &blk)
  end
end