Repository: OpenMined/PySyft

Issue:
Test worker serialization
**Is your feature request related to a problem? Please describe.**

We should have tests for Virtual Workers serialization since this is crucial for [Grid](https://github.com/OpenMined/Grid).


Discussion:

Resolving patch:
diff --git a/syft/serde/serde.py b/syft/serde/serde.py
--- a/syft/serde/serde.py
+++ b/syft/serde/serde.py
@@ -69,10 +69,7 @@
     + list(MAP_TORCH_SIMPLIFIERS_AND_DETAILERS.items())
 )
 
-# Objects that we can run force full simplilfy on
-MAP_TO_FORCE_FULL_SIMPLIFY = [VirtualWorker]
-
-# If a object implements its own simplifer and detailer function it should be stored in this list
+# If a object implements its own simplify and detail functions it should be stored in this list
 OBJ_SIMPLIFIER_AND_DETAILERS = [
     AdditiveSharingTensor,
     LoggingTensor,
@@ -84,6 +81,10 @@
     VirtualWorker,
 ]
 
+# If a object implements its own force_simplify and force_detail functions it should be stored in this list
+OBJ_FORCE_FULL_SIMPLIFIER_AND_DETAILERS = [VirtualWorker]
+
+
 EXCEPTION_SIMPLIFIER_AND_DETAILERS = [GetNotPermittedError, ResponseSignatureError]
 
 # COMPRESSION SCHEME INT CODES
@@ -96,11 +97,9 @@ def _force_full_simplify(obj: object) -> object:
     current_type = type(obj)
 
     if current_type in forced_full_simplifiers:
-
         left = forced_full_simplifiers[current_type][0]
 
         right = forced_full_simplifiers[current_type][1]
-
         right = right(obj)
 
         result = (left, right)
@@ -110,22 +109,6 @@ def _force_full_simplify(obj: object) -> object:
     return result
 
 
-def _force_full_detail(worker: AbstractWorker, worker_tuple: tuple) -> tuple:
-    worker_id, _objects, auto_add = worker_tuple
-    worker_id = _detail(worker, worker_id)
-
-    result = sy.VirtualWorker(sy.hook, worker_id, auto_add=auto_add)
-    _objects = _detail(worker, _objects)
-    result._objects = _objects
-
-    # make sure they weren't accidentally double registered
-    for _, obj in _objects.items():
-        if obj.id in worker._objects:
-            del worker._objects[obj.id]
-
-    return result
-
-
 ## SECTION: dinamically generate simplifiers and detailers
 def _generate_simplifiers_and_detailers():
     """Generate simplifiers, forced full simplifiers and detailers."""
@@ -155,17 +138,17 @@ def _add_simplifier_and_detailer(curr_type, simplifier, detailer, forced=False):
         simplifier, detailer = syft_type.simplify, syft_type.detail
         _add_simplifier_and_detailer(syft_type, simplifier, detailer)
 
-    # Add forced full simplifiers
-    for curr_type in MAP_TO_FORCE_FULL_SIMPLIFY:
-        _add_simplifier_and_detailer(
-            curr_type, _force_full_simplify, _force_full_detail, forced=True
-        )
+    # Register syft objects with custom force_simplify and force_detail methods
+    for syft_type in OBJ_FORCE_FULL_SIMPLIFIER_AND_DETAILERS:
+        force_simplifier, force_detailer = syft_type.force_simplify, syft_type.force_detail
+        _add_simplifier_and_detailer(syft_type, force_simplifier, force_detailer, forced=True)
 
     return simplifiers, forced_full_simplifiers, detailers
 
 
 simplifiers, forced_full_simplifiers, detailers = _generate_simplifiers_and_detailers()
 
+
 ## SECTION:  High Level Public Functions (these are the ones you use)
 def serialize(
     obj: object,
diff --git a/syft/workers/virtual.py b/syft/workers/virtual.py
--- a/syft/workers/virtual.py
+++ b/syft/workers/virtual.py
@@ -13,11 +13,27 @@ def _recv_msg(self, message: bin) -> bin:
         return self.recv_msg(message)
 
     @staticmethod
-    def simplify(worker: AbstractWorker) -> tuple:
-        """
+    def force_simplify(worker: AbstractWorker) -> tuple:
+        return (sy.serde._simplify(worker.id), sy.serde._simplify(worker._objects), worker.auto_add)
 
-        """
+    @staticmethod
+    def force_detail(worker: AbstractWorker, worker_tuple: tuple) -> tuple:
+        worker_id, _objects, auto_add = worker_tuple
+        worker_id = sy.serde._detail(worker, worker_id)
+
+        result = sy.VirtualWorker(sy.hook, worker_id, auto_add=auto_add)
+        _objects = sy.serde._detail(worker, _objects)
+        result._objects = _objects
 
+        # make sure they weren't accidentally double registered
+        for _, obj in _objects.items():
+            if obj.id in worker._objects:
+                del worker._objects[obj.id]
+
+        return result
+
+    @staticmethod
+    def simplify(worker: AbstractWorker) -> tuple:
         return (sy.serde._simplify(worker.id),)
 
     @staticmethod


Test patch:
diff --git a/test/test_serde.py b/test/test_serde.py
--- a/test/test_serde.py
+++ b/test/test_serde.py
@@ -553,8 +553,6 @@ def test_pointer_tensor(hook, workers):
     )
     t_serialized = serde.serialize(t)
     t_serialized_deserialized = serde.deserialize(t_serialized)
-    print(f"t.location - {t.location}")
-    print(f"t_serialized_deserialized.location - {t_serialized_deserialized.location}")
     assert t.id == t_serialized_deserialized.id
     assert t.location.id == t_serialized_deserialized.location.id
     assert t.id_at_location == t_serialized_deserialized.id_at_location
@@ -643,6 +641,35 @@ def foo(x):
     assert foo.code == foo_received.code
 
 
+def test_serde_virtual_worker(hook):
+    virtual_worker = syft.VirtualWorker(hook=hook, id="deserialized_worker1")
+    # Populate worker
+    tensor1, tensor2 = torch.tensor([1.0, 2.0]), torch.tensor([0.0])
+    ptr1, ptr2 = tensor1.send(virtual_worker), tensor2.send(virtual_worker)
+
+    serialized_worker = serde.serialize(virtual_worker, force_full_simplification=False)
+    deserialized_worker = serde.deserialize(serialized_worker)
+
+    assert virtual_worker.id == deserialized_worker.id
+
+
+def test_full_serde_virtual_worker(hook):
+    virtual_worker = syft.VirtualWorker(hook=hook, id="deserialized_worker2")
+    # Populate worker
+    tensor1, tensor2 = torch.tensor([1.0, 2.0]), torch.tensor([0.0])
+    ptr1, ptr2 = tensor1.send(virtual_worker), tensor2.send(virtual_worker)
+
+    serialized_worker = serde.serialize(virtual_worker, force_full_simplification=True)
+
+    deserialized_worker = serde.deserialize(serialized_worker)
+
+    assert virtual_worker.id == deserialized_worker.id
+    assert virtual_worker.auto_add == deserialized_worker.auto_add
+    assert len(deserialized_worker._objects) == 2
+    assert tensor1.id in deserialized_worker._objects
+    assert tensor2.id in deserialized_worker._objects
+
+
 def test_serde_object_wrapper_traced_module():
 
     data = torch.tensor([[-1, 2.0], [0, 1.1], [-1, 2.1], [0, 1.2]])