From 80ffb1f9c15cdccf2f25652acfec21cc28f186f3 Mon Sep 17 00:00:00 2001
From: Jonathan del Strother <jdelStrother@gmail.com>
Date: Mon, 25 Nov 2019 11:10:55 +0000
Subject: [PATCH] Add support for dump/restore between MockRedis instances

---
 lib/mock_redis/database.rb    | 16 ++++++++++++
 spec/commands/dump_spec.rb    | 19 ++++++++++++++
 spec/commands/restore_spec.rb | 47 +++++++++++++++++++++++++++++++++++
 3 files changed, 82 insertions(+)
 create mode 100644 spec/commands/dump_spec.rb
 create mode 100644 spec/commands/restore_spec.rb

diff --git a/lib/mock_redis/database.rb b/lib/mock_redis/database.rb
index d006f812..383d4abe 100644
--- a/lib/mock_redis/database.rb
+++ b/lib/mock_redis/database.rb
@@ -124,6 +124,22 @@ def flushdb
       'OK'
     end
 
+    def dump(key)
+      value = data[key]
+      value ? Marshal.dump(value) : nil
+    end
+
+    def restore(key, ttl, value, replace: false)
+      if !replace && exists(key)
+        raise Redis::CommandError, 'BUSYKEY Target key name already exists.'
+      end
+      data[key] = Marshal.load(value) # rubocop:disable Security/MarshalLoad
+      if ttl > 0
+        pexpire(key, ttl)
+      end
+      'OK'
+    end
+
     def keys(format = '*')
       data.keys.grep(redis_pattern_to_ruby_regex(format))
     end
diff --git a/spec/commands/dump_spec.rb b/spec/commands/dump_spec.rb
new file mode 100644
index 00000000..1ba1177a
--- /dev/null
+++ b/spec/commands/dump_spec.rb
@@ -0,0 +1,19 @@
+require 'spec_helper'
+
+describe '#dump(key)' do
+  before do
+    @key = 'mock-redis-test:45794'
+    # These are mock-only, since our dump/restore implementations
+    # aren't compatible with real redis.
+    @mock = @redises.mock
+  end
+
+  it 'returns nil for keys that do not exist' do
+    @mock.dump(@key).should be_nil
+  end
+
+  it 'returns a serialized value for keys that do exist' do
+    @mock.set(@key, '2')
+    @mock.dump(@key).should == Marshal.dump('2')
+  end
+end
diff --git a/spec/commands/restore_spec.rb b/spec/commands/restore_spec.rb
new file mode 100644
index 00000000..c0f00604
--- /dev/null
+++ b/spec/commands/restore_spec.rb
@@ -0,0 +1,47 @@
+require 'spec_helper'
+
+describe '#restore(key, ttl, value)' do
+  before do
+    @key = 'mock-redis-test:45794'
+    @src = MockRedis.new
+    @src.set(@key, '123')
+    @dumped_value = @src.dump(@key)
+    @dst = MockRedis.new
+    @now = Time.now.round
+    Time.stub(:now).and_return(@now)
+  end
+
+  it 'allows dump/restoring values between two redis instances' do
+    expect(@dst.restore(@key, 0, @dumped_value)).to eq('OK')
+    expect(@dst.get(@key)).to eq('123')
+    expect(@dst.pttl(@key)).to eq(-1)
+  end
+
+  context 'when the key being restored to already exists' do
+    before do
+      @dst.set(@key, '456')
+    end
+
+    it 'raises an error by default' do
+      expect { @dst.restore(@key, 0, @dumped_value) }.to raise_error(Redis::CommandError)
+      expect(@dst.get(@key)).to eq('456')
+    end
+
+    it 'allows replacing the key if replace==true' do
+      expect(@dst.restore(@key, 0, @dumped_value, replace: true)).to eq('OK')
+      expect(@dst.get(@key)).to eq('123')
+    end
+  end
+
+  it 'sets ttl in ms' do
+    @dst.restore(@key, 500, @dumped_value)
+    expect(@dst.pttl(@key)).to eq(500)
+  end
+
+  it 'can dump/restore more complex data types' do
+    key = 'a_hash'
+    @src.mapped_hmset(key, foo: 'bar')
+    @dst.restore(key, 0, @src.dump(key))
+    expect(@dst.hgetall(key)).to eq('foo' => 'bar')
+  end
+end