Added Custom Ref Handlers

tests/test_serializer.py

@@ -6,7 +6,8 @@ from enum import Enum
 
 import pytest
 
-from dbengine.serializer import TAG_TUPLE, TAG_SET, Serializer, TAG_OBJECT, TAG_ID, TAG_REF
+from dbengine.serializer import TAG_TUPLE, TAG_SET, Serializer, TAG_OBJECT, TAG_ID, TAG_REF, TAG_DIGEST
+from dbengine.handlers import BaseRefHandler, handlers
 
 
 class Obj:
@@ -72,28 +73,74 @@ class DummyComplexClass:
   prop3: ObjEnum
 
 
+class BytesData:
+  """Simple class to hold binary data for testing BaseRefHandler"""
+
+  def __init__(self, data: bytes):
+    self.data = data
+
+  def __eq__(self, other):
+    if not isinstance(other, BytesData):
+      return False
+    return self.data == other.data
+
+  def __hash__(self):
+    return hash(self.data)
+
+
+class BytesDataHandler(BaseRefHandler):
+  """Test handler that stores BytesData using refs/ directory"""
+
+  def is_eligible_for(self, obj):
+    return isinstance(obj, BytesData)
+
+  def tag(self):
+    return "BytesData"
+
+  def serialize_to_bytes(self, obj) -> bytes:
+    """Just return the raw bytes from the object"""
+    return obj.data
+
+  def deserialize_from_bytes(self, data: bytes) -> object:
+    """Reconstruct BytesData from raw bytes"""
+    return BytesData(data)
+
+
 class DummyRefHelper:
   """
   When something is too complicated to serialize, we just default to pickle
   That is what this helper class is doing
   """
-  
+
   def __init__(self):
     self.refs = {}
-  
+
   def save_ref(self, obj):
     sha256_hash = hashlib.sha256()
-    
+
     pickled_data = pickle.dumps(obj)
     sha256_hash.update(pickled_data)
     digest = sha256_hash.hexdigest()
-    
+
     self.refs[digest] = pickled_data
     return digest
-  
+
   def load_ref(self, digest):
     return pickle.loads(self.refs[digest])
 
+  def save_ref_from_bytes(self, data: bytes):
+    """Save raw bytes for BaseRefHandler"""
+    sha256_hash = hashlib.sha256()
+    sha256_hash.update(data)
+    digest = sha256_hash.hexdigest()
+
+    self.refs[digest] = data
+    return digest
+
+  def load_ref_to_bytes(self, digest):
+    """Load raw bytes for BaseRefHandler"""
+    return self.refs[digest]
+
 
 @pytest.mark.parametrize("obj, expected", [
     (1, 1),
@@ -260,9 +307,81 @@ def test_can_use_refs_when_circular_reference():
 def test_i_can_serialize_date():
   obj = datetime.date.today()
   serializer = Serializer()
-  
+
   flatten = serializer.serialize(obj)
-  
+
   decoded = serializer.deserialize(flatten)
-  
+
   assert decoded == obj
+
+
+def test_i_can_serialize_with_base_ref_handler():
+  """Test that BaseRefHandler correctly stores data in refs/ with TAG_DIGEST"""
+  # Register the test handler
+  test_handler = BytesDataHandler()
+  handlers.register_handler(test_handler)
+
+  try:
+    # Create test data
+    test_bytes = b"Hello, this is binary data for testing!"
+    obj = BytesData(test_bytes)
+
+    # Serialize with ref_helper
+    ref_helper = DummyRefHelper()
+    serializer = Serializer(ref_helper)
+
+    flatten = serializer.serialize(obj)
+
+    # Verify structure: should have TAG_SPECIAL and TAG_DIGEST
+    assert "__special__" in flatten
+    assert flatten["__special__"] == "BytesData"
+    assert TAG_DIGEST in flatten
+
+    # Verify data was stored in refs
+    digest = flatten[TAG_DIGEST]
+    assert digest in ref_helper.refs
+    assert ref_helper.refs[digest] == test_bytes
+
+    # Deserialize and verify
+    decoded = serializer.deserialize(flatten)
+    assert decoded == obj
+    assert decoded.data == test_bytes
+
+  finally:
+    # Clean up: remove the test handler
+    handlers.handlers.remove(test_handler)
+
+
+def test_i_can_serialize_object_containing_base_ref_handler_data():
+  """Test that objects containing BaseRefHandler-managed data work correctly"""
+  # Register the test handler
+  test_handler = BytesDataHandler()
+  handlers.register_handler(test_handler)
+
+  try:
+    # Create an object that contains BytesData
+    bytes_obj = BytesData(b"Some binary content")
+    wrapper = Obj(1, "test", bytes_obj)
+
+    # Serialize
+    ref_helper = DummyRefHelper()
+    serializer = Serializer(ref_helper)
+
+    flatten = serializer.serialize(wrapper)
+
+    # Verify structure
+    assert flatten[TAG_OBJECT] == 'tests.test_serializer.Obj'
+    assert flatten['a'] == 1
+    assert flatten['b'] == "test"
+    assert flatten['c']["__special__"] == "BytesData"
+    assert TAG_DIGEST in flatten['c']
+
+    # Deserialize and verify
+    decoded = serializer.deserialize(flatten)
+    assert decoded.a == wrapper.a
+    assert decoded.b == wrapper.b
+    assert decoded.c == wrapper.c
+
+  finally:
+    # Clean up
+    handlers.handlers.remove(test_handler)