From 2915fb9d85bdf88b115e09a6f6c17d17282ae068 Mon Sep 17 00:00:00 2001
From: Jean byroot Boussier <jean.boussier+github@shopify.com>
Date: Thu, 8 Apr 2021 19:41:08 +0200
Subject: [PATCH] Allow to configure the codec (#195)

Co-authored-by: Jean Boussier <jean.boussier@gmail.com>
---
 lib/memcached/client.rb  |  2 +-
 test/unit/client_test.rb | 25 +++++++++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/lib/memcached/client.rb b/lib/memcached/client.rb
index f0cb17a..704bc38 100644
--- a/lib/memcached/client.rb
+++ b/lib/memcached/client.rb
@@ -23,7 +23,7 @@ module Memcached
     def initialize(config = "localhost:11211", options = {})
       options = DEFAULTS.merge(options)
 
-      @codec = Memcached::MarshalCodec
+      @codec = options.delete(:codec) || Memcached::MarshalCodec
       @default_ttl = options.delete(:ttl) || 0
       @prefix = options.delete(:prefix_key)
       @credentials = options.delete(:credentials)
diff --git a/test/unit/client_test.rb b/test/unit/client_test.rb
index 66cc86f..0960070 100644
--- a/test/unit/client_test.rb
+++ b/test/unit/client_test.rb
@@ -175,6 +175,31 @@ class ClientTest < BaseTest
     assert(failed < keys.size / 3, "#{failed} failed out of #{keys.size}")
   end
 
+  module YAMLCodec
+    extend self
+
+    FLAG = Memcached::Client::FLAG_ENCODED
+
+    def encode(key, value, flags)
+      [ YAML.dump(value), flags | FLAG ]
+    end
+
+    def decode(key, value, flags)
+      if (flags & FLAG) != 0
+        YAML.load(value)
+      else
+        value
+      end
+    end
+  end
+
+  def test_custom_codec
+    cache = Memcached::Client.new(@servers, codec: YAMLCodec)
+    cache.set("yaml", :symbol)
+    assert_equal :symbol, cache.get("yaml")
+    assert_equal YAML.dump(:symbol), cache.get("yaml", raw: true)
+  end
+
   # Touch command
 
   def test_touch_missing_key