Add infrastructure for capturing Redis commands in tests (#304)

redis-rb · Dec 13, 2023 · 11d4b68 · 11d4b68
1 parent 8a3d03e
commit 11d4b68
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 5 deletions.
diff --git a/test/command_capture_middleware.rb b/test/command_capture_middleware.rb
@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module CommandCaptureMiddleware
+  CapturedCommand = Struct.new(:server_url, :command, :pipelined, keyword_init: true) do
+    def inspect
+      "#<#{self.class.name} [on #{server_url}] #{command.join(' ')} >"
+    end
+  end
+
+  def call(command, redis_config)
+    redis_config.custom[:captured_commands] << CapturedCommand.new(
+      server_url: redis_config.server_url,
+      command: command,
+      pipelined: false
+    )
+    super
+  end
+
+  def call_pipelined(commands, redis_config)
+    commands.map do |command|
+      redis_config.custom[:captured_commands] << CapturedCommand.new(
+        server_url: redis_config.server_url,
+        command: command,
+        pipelined: true
+      )
+    end
+    super
+  end
+end
diff --git a/test/redis_client/test_cluster.rb b/test/redis_client/test_cluster.rb
@@ -6,9 +6,11 @@ class RedisClient
   class TestCluster
     module Mixin
       def setup
+        @captured_commands = []
         @client = new_test_client
         @client.call('FLUSHDB')
         wait_for_replication
+        @captured_commands.clear
       end
 
       def teardown
@@ -516,10 +518,12 @@ def hiredis_used?
     class PrimaryOnly < TestingWrapper
       include Mixin
 
-      def new_test_client
+      def new_test_client(capture_buffer: @captured_commands)
         config = ::RedisClient::ClusterConfig.new(
           nodes: TEST_NODE_URIS,
           fixed_hostname: TEST_FIXED_HOSTNAME,
+          middlewares: [CommandCaptureMiddleware],
+          custom: { captured_commands: capture_buffer },
           **TEST_GENERIC_OPTIONS
         )
         ::RedisClient::Cluster.new(config)
@@ -529,12 +533,14 @@ def new_test_client
     class ScaleReadRandom < TestingWrapper
       include Mixin
 
-      def new_test_client
+      def new_test_client(capture_buffer: @captured_commands)
         config = ::RedisClient::ClusterConfig.new(
           nodes: TEST_NODE_URIS,
           replica: true,
           replica_affinity: :random,
           fixed_hostname: TEST_FIXED_HOSTNAME,
+          middlewares: [CommandCaptureMiddleware],
+          custom: { captured_commands: capture_buffer },
           **TEST_GENERIC_OPTIONS
         )
         ::RedisClient::Cluster.new(config)
@@ -544,12 +550,14 @@ def new_test_client
     class ScaleReadRandomWithPrimary < TestingWrapper
       include Mixin
 
-      def new_test_client
+      def new_test_client(capture_buffer: @captured_commands)
         config = ::RedisClient::ClusterConfig.new(
           nodes: TEST_NODE_URIS,
           replica: true,
           replica_affinity: :random_with_primary,
           fixed_hostname: TEST_FIXED_HOSTNAME,
+          middlewares: [CommandCaptureMiddleware],
+          custom: { captured_commands: capture_buffer },
           **TEST_GENERIC_OPTIONS
         )
         ::RedisClient::Cluster.new(config)
@@ -559,12 +567,14 @@ def new_test_client
     class ScaleReadLatency < TestingWrapper
       include Mixin
 
-      def new_test_client
+      def new_test_client(capture_buffer: @captured_commands)
         config = ::RedisClient::ClusterConfig.new(
           nodes: TEST_NODE_URIS,
           replica: true,
           replica_affinity: :latency,
           fixed_hostname: TEST_FIXED_HOSTNAME,
+          middlewares: [CommandCaptureMiddleware],
+          custom: { captured_commands: capture_buffer },
           **TEST_GENERIC_OPTIONS
         )
         ::RedisClient::Cluster.new(config)
@@ -574,10 +584,12 @@ def new_test_client
     class Pooled < TestingWrapper
       include Mixin
 
-      def new_test_client
+      def new_test_client(capture_buffer: @captured_commands)
         config = ::RedisClient::ClusterConfig.new(
           nodes: TEST_NODE_URIS,
           fixed_hostname: TEST_FIXED_HOSTNAME,
+          middlewares: [CommandCaptureMiddleware],
+          custom: { captured_commands: capture_buffer },
           **TEST_GENERIC_OPTIONS
         )
         ::RedisClient::Cluster.new(config, pool: { timeout: TEST_TIMEOUT_SEC, size: 2 })

diff --git a/test/testing_helper.rb b/test/testing_helper.rb
@@ -6,6 +6,7 @@
 require 'redis-cluster-client'
 require 'testing_constants'
 require 'cluster_controller'
+require 'command_capture_middleware'
 
 case ENV.fetch('REDIS_CONNECTION_DRIVER', 'ruby')
 when 'hiredis' then require 'hiredis-client'