Repository: OpenMined/PySyft

Issue:
Allow worker registration of objects within python context
**Is your feature request related to a problem? Please describe.**
Tons of test and code samples using plans are currently like this:

```python
hook.local_worker.is_client_worker = False
# my awesome code
hook.local_worker.is_client_worker = True
```


**Describe the solution you'd like**
```python
with self.owner.registration_enabled():
    # my awesome code
```

**Describe alternatives you've considered**
Enabling registration all the time :O




Discussion:
I can start looking at this one

Resolving patch:
diff --git a/syft/workers/base.py b/syft/workers/base.py
--- a/syft/workers/base.py
+++ b/syft/workers/base.py
@@ -1,4 +1,6 @@
 from abc import abstractmethod
+from contextlib import contextmanager
+
 import logging
 from typing import Callable
 from typing import List
@@ -202,6 +204,14 @@ def _recv_msg(self, message: bin):
         """
         raise NotImplementedError  # pragma: no cover
 
+    @contextmanager
+    def registration_enabled(self):
+        self.is_client_worker = False
+        try:
+            yield self
+        finally:
+            self.is_client_worker = True
+
     def remove_worker_from_registry(self, worker_id):
         """Removes a worker from the dictionary of known workers.
         Args:


Test patch:
diff --git a/test/message/test_plan.py b/test/message/test_plan.py
--- a/test/message/test_plan.py
+++ b/test/message/test_plan.py
@@ -27,24 +27,22 @@ def plan_abs(data):
 
 def test_stateful_plan_built_automatically(hook):
 
-    hook.local_worker.is_client_worker = False
+    with hook.local_worker.registration_enabled():
 
-    @sy.func2plan(args_shape=[(1,)], state=(th.tensor([1.0]),))
-    def foo(x, state):
-        bias, = state.read()
-        x = x * 2
-        return x + bias
-
-    assert isinstance(foo.__str__(), str)
-    assert len(foo.readable_plan) > 0
-    assert foo.is_built
+        @sy.func2plan(args_shape=[(1,)], state=(th.tensor([1.0]),))
+        def foo(x, state):
+            bias, = state.read()
+            x = x * 2
+            return x + bias
 
-    t = th.tensor([1.0, 2])
-    x = foo(t)
+        assert isinstance(foo.__str__(), str)
+        assert len(foo.readable_plan) > 0
+        assert foo.is_built
 
-    assert (x == th.tensor([3.0, 5])).all()
+        t = th.tensor([1.0, 2])
+        x = foo(t)
 
-    hook.local_worker.is_client_worker = True
+        assert (x == th.tensor([3.0, 5])).all()
 
 
 def test_plan_build():
@@ -62,20 +60,18 @@ def plan_abs(data):
 
 
 def test_stateful_plan_build(hook):
-    hook.local_worker.is_client_worker = False
+    with hook.local_worker.registration_enabled():
 
-    @sy.func2plan(state=(th.tensor([1.0]),))
-    def foo(x, state):
-        bias, = state.read()
-        x = x * 2
-        return x + bias
-
-    t = th.tensor([1.0, 2])
-    x = foo(t)
+        @sy.func2plan(state=(th.tensor([1.0]),))
+        def foo(x, state):
+            bias, = state.read()
+            x = x * 2
+            return x + bias
 
-    assert (x == th.tensor([3.0, 5])).all()
+        t = th.tensor([1.0, 2])
+        x = foo(t)
 
-    hook.local_worker.is_client_worker = True
+        assert (x == th.tensor([3.0, 5])).all()
 
 
 def test_plan_built_automatically_with_any_dimension():
@@ -136,59 +132,55 @@ def forward(self, x):
 
 def test_plan_method_execute_locally(hook):
 
-    hook.local_worker.is_client_worker = False
-
-    class Net(sy.Plan):
-        def __init__(self):
-            super(Net, self).__init__()
-            self.fc1 = nn.Linear(2, 3)
-            self.fc2 = nn.Linear(3, 2)
-            self.fc3 = nn.Linear(2, 1)
+    with hook.local_worker.registration_enabled():
 
-        def forward(self, x):
-            x = F.relu(self.fc1(x))
-            x = self.fc2(x)
-            x = self.fc3(x)
-            return F.log_softmax(x, dim=0)
+        class Net(sy.Plan):
+            def __init__(self):
+                super(Net, self).__init__()
+                self.fc1 = nn.Linear(2, 3)
+                self.fc2 = nn.Linear(3, 2)
+                self.fc3 = nn.Linear(2, 1)
 
-    model = Net()
+            def forward(self, x):
+                x = F.relu(self.fc1(x))
+                x = self.fc2(x)
+                x = self.fc3(x)
+                return F.log_softmax(x, dim=0)
 
-    model.build(th.tensor([1.0, 2]))
+        model = Net()
 
-    # Call one time
-    assert model(th.tensor([1.0, 2])) == 0
+        model.build(th.tensor([1.0, 2]))
 
-    # Call one more time
-    assert model(th.tensor([1.0, 2.1])) == 0
+        # Call one time
+        assert model(th.tensor([1.0, 2])) == 0
 
-    hook.local_worker.is_client_worker = True
+        # Call one more time
+        assert model(th.tensor([1.0, 2.1])) == 0
 
 
 def test_stateful_plan_method_execute_locally(hook):
 
-    hook.local_worker.is_client_worker = False
+    with hook.local_worker.registration_enabled():
 
-    class Net(sy.Plan):
-        def __init__(self):
-            super(Net, self).__init__()
-            self.fc1 = nn.Linear(2, 1)
-            self.bias = th.tensor([1000.0])
-
-        def forward(self, x):
-            x = self.fc1(x)
-            return F.log_softmax(x, dim=0) + self.bias
+        class Net(sy.Plan):
+            def __init__(self):
+                super(Net, self).__init__()
+                self.fc1 = nn.Linear(2, 1)
+                self.bias = th.tensor([1000.0])
 
-    model = Net()
+            def forward(self, x):
+                x = self.fc1(x)
+                return F.log_softmax(x, dim=0) + self.bias
 
-    model.build(th.tensor([1.0, 2]))
+        model = Net()
 
-    # Call one time
-    assert model(th.tensor([1.0, 2])) == th.tensor([1000.0])
+        model.build(th.tensor([1.0, 2]))
 
-    # Call one more time
-    assert model(th.tensor([1.0, 2.1])) == th.tensor([1000.0])
+        # Call one time
+        assert model(th.tensor([1.0, 2])) == th.tensor([1000.0])
 
-    hook.local_worker.is_client_worker = True
+        # Call one more time
+        assert model(th.tensor([1.0, 2.1])) == th.tensor([1000.0])
 
 
 def test_plan_multiple_send(workers):
@@ -217,30 +209,28 @@ def plan_abs(data):
 def test_stateful_plan_multiple_send(hook, workers):
     me, bob, alice = workers["me"], workers["bob"], workers["alice"]
 
-    hook.local_worker.is_client_worker = False
+    with hook.local_worker.registration_enabled():
 
-    @sy.func2plan(args_shape=[(1,)], state=(th.tensor([1.0]),))
-    def plan_abs(x, state):
-        bias, = state.read()
-        x = x.abs()
-        return x + bias
-
-    plan_ptr = plan_abs.send(bob)
-    x_ptr = th.tensor([-1.0, 7, 3]).send(bob)
-    p = plan_ptr(x_ptr)
-    res = p.get()
+        @sy.func2plan(args_shape=[(1,)], state=(th.tensor([1.0]),))
+        def plan_abs(x, state):
+            bias, = state.read()
+            x = x.abs()
+            return x + bias
 
-    assert (res == th.tensor([2.0, 8, 4])).all()
+        plan_ptr = plan_abs.send(bob)
+        x_ptr = th.tensor([-1.0, 7, 3]).send(bob)
+        p = plan_ptr(x_ptr)
+        res = p.get()
 
-    # Test get / send plan
-    plan_ptr = plan_abs.send(alice)
+        assert (res == th.tensor([2.0, 8, 4])).all()
 
-    x_ptr = th.tensor([-1.0, 2, 3]).send(alice)
-    p = plan_ptr(x_ptr)
-    res = p.get()
-    assert (res == th.tensor([2.0, 3, 4])).all()
+        # Test get / send plan
+        plan_ptr = plan_abs.send(alice)
 
-    hook.local_worker.is_client_worker = True
+        x_ptr = th.tensor([-1.0, 2, 3]).send(alice)
+        p = plan_ptr(x_ptr)
+        res = p.get()
+        assert (res == th.tensor([2.0, 3, 4])).all()
 
 
 def test_plan_built_on_class(hook):
@@ -248,50 +238,48 @@ def test_plan_built_on_class(hook):
     Test class Plans and plan send / get / send
     """
 
-    hook.local_worker.is_client_worker = False
+    with hook.local_worker.registration_enabled():
 
-    x11 = th.tensor([-1, 2.0]).tag("input_data")
-    x21 = th.tensor([-1, 2.0]).tag("input_data")
+        x11 = th.tensor([-1, 2.0]).tag("input_data")
+        x21 = th.tensor([-1, 2.0]).tag("input_data")
 
-    device_1 = sy.VirtualWorker(hook, id="device_1", data=(x11,))
-    device_2 = sy.VirtualWorker(hook, id="device_2", data=(x21,))
+        device_1 = sy.VirtualWorker(hook, id="device_1", data=(x11,))
+        device_2 = sy.VirtualWorker(hook, id="device_2", data=(x21,))
 
-    class Net(sy.Plan):
-        def __init__(self):
-            super(Net, self).__init__()
-            self.fc1 = nn.Linear(2, 3)
-            self.fc2 = nn.Linear(3, 1)
-
-            self.bias = th.tensor([1000.0])
+        class Net(sy.Plan):
+            def __init__(self):
+                super(Net, self).__init__()
+                self.fc1 = nn.Linear(2, 3)
+                self.fc2 = nn.Linear(3, 1)
 
-        def forward(self, x):
-            x = F.relu(self.fc1(x))
-            x = self.fc2(x)
-            return F.log_softmax(x, dim=0) + self.bias
+                self.bias = th.tensor([1000.0])
 
-    net = Net()
+            def forward(self, x):
+                x = F.relu(self.fc1(x))
+                x = self.fc2(x)
+                return F.log_softmax(x, dim=0) + self.bias
 
-    # build
-    net.build(th.tensor([1, 2.0]))
+        net = Net()
 
-    net_ptr = net.send(device_1)
-    pointer_to_data = device_1.search("input_data")[0]
-    pointer_to_result = net_ptr(pointer_to_data)
+        # build
+        net.build(th.tensor([1, 2.0]))
 
-    result = pointer_to_result.get()
-    assert isinstance(result, th.Tensor)
-    assert result == th.tensor([1000.0])
+        net_ptr = net.send(device_1)
+        pointer_to_data = device_1.search("input_data")[0]
+        pointer_to_result = net_ptr(pointer_to_data)
 
-    net_ptr = net.send(device_2)
+        result = pointer_to_result.get()
+        assert isinstance(result, th.Tensor)
+        assert result == th.tensor([1000.0])
 
-    pointer_to_data = device_2.search("input_data")[0]
-    pointer_to_result = net_ptr(pointer_to_data)
+        net_ptr = net.send(device_2)
 
-    result = pointer_to_result.get()
-    assert isinstance(result, th.Tensor)
-    assert result == th.tensor([1000.0])
+        pointer_to_data = device_2.search("input_data")[0]
+        pointer_to_result = net_ptr(pointer_to_data)
 
-    hook.local_worker.is_client_worker = True
+        result = pointer_to_result.get()
+        assert isinstance(result, th.Tensor)
+        assert result == th.tensor([1000.0])
 
 
 def test_multiple_workers(workers):
@@ -316,367 +304,348 @@ def plan_abs(data):
 def test_stateful_plan_multiple_workers(hook, workers):
     me, bob, alice = workers["me"], workers["bob"], workers["alice"]
 
-    hook.local_worker.is_client_worker = False
-
-    @sy.func2plan(args_shape=[(1,)], state=(th.tensor([1]),))
-    def plan_abs(x, state):
-        bias, = state.read()
-        x = x.abs()
-        return x + bias
+    with hook.local_worker.registration_enabled():
 
-    plan_ptr = plan_abs.send(bob, alice)
-    x_ptr = th.tensor([-1, 7, 3]).send(bob)
-    p = plan_ptr(x_ptr)
-    x_abs = p.get()
-    assert (x_abs == th.tensor([2, 8, 4])).all()
+        @sy.func2plan(args_shape=[(1,)], state=(th.tensor([1]),))
+        def plan_abs(x, state):
+            bias, = state.read()
+            x = x.abs()
+            return x + bias
 
-    x_ptr = th.tensor([-1, 9, 3]).send(alice)
-    p = plan_ptr(x_ptr)
-    x_abs = p.get()
-    assert (x_abs == th.tensor([2, 10, 4])).all()
+        plan_ptr = plan_abs.send(bob, alice)
+        x_ptr = th.tensor([-1, 7, 3]).send(bob)
+        p = plan_ptr(x_ptr)
+        x_abs = p.get()
+        assert (x_abs == th.tensor([2, 8, 4])).all()
 
-    hook.local_worker.is_client_worker = True
+        x_ptr = th.tensor([-1, 9, 3]).send(alice)
+        p = plan_ptr(x_ptr)
+        x_abs = p.get()
+        assert (x_abs == th.tensor([2, 10, 4])).all()
 
 
 def test_fetch_plan(hook, workers):
     alice = workers["alice"]
 
-    hook.local_worker.is_client_worker = False
-
-    @sy.func2plan(args_shape=[(1,)])
-    def plan(data):
-        return data * 3
+    with hook.local_worker.registration_enabled():
 
-    plan.send(alice)
+        @sy.func2plan(args_shape=[(1,)])
+        def plan(data):
+            return data * 3
 
-    # Fetch plan
-    fetched_plan = plan.owner.fetch_plan(plan.id, alice)
+        plan.send(alice)
 
-    # Execute it locally
-    x = th.tensor([-1.0, 2, 3])
-    assert (plan(x) == th.tensor([-3.0, 6, 9])).all()
-    assert (fetched_plan(x) == th.tensor([-3.0, 6, 9])).all()
-    assert fetched_plan.forward is None
-    assert fetched_plan.is_built
+        # Fetch plan
+        fetched_plan = plan.owner.fetch_plan(plan.id, alice)
 
-    hook.local_worker.is_client_worker = True
+        # Execute it locally
+        x = th.tensor([-1.0, 2, 3])
+        assert (plan(x) == th.tensor([-3.0, 6, 9])).all()
+        assert (fetched_plan(x) == th.tensor([-3.0, 6, 9])).all()
+        assert fetched_plan.forward is None
+        assert fetched_plan.is_built
 
 
 @pytest.mark.parametrize("is_func2plan", [True, False])
 def test_fetch_stateful_plan(hook, is_func2plan, workers):
-    hook.local_worker.is_client_worker = False
 
-    if is_func2plan:
+    with hook.local_worker.registration_enabled():
+        if is_func2plan:
 
-        @sy.func2plan(args_shape=[(1,)], state=(th.tensor([1.0]),))
-        def plan(data, state):
-            bias, = state.read()
-            return data * bias
+            @sy.func2plan(args_shape=[(1,)], state=(th.tensor([1.0]),))
+            def plan(data, state):
+                bias, = state.read()
+                return data * bias
 
-    else:
+        else:
 
-        class Net(sy.Plan):
-            def __init__(self):
-                super(Net, self).__init__()
-                self.fc1 = nn.Linear(1, 1)
+            class Net(sy.Plan):
+                def __init__(self):
+                    super(Net, self).__init__()
+                    self.fc1 = nn.Linear(1, 1)
 
-            def forward(self, x):
-                return self.fc1(x)
+                def forward(self, x):
+                    return self.fc1(x)
 
-        plan = Net()
-        plan.build(th.tensor([1.2]))
-
-    alice = workers["alice"]
-    plan_ptr = plan.send(alice)
+            plan = Net()
+            plan.build(th.tensor([1.2]))
 
-    # Fetch plan
-    fetched_plan = plan.owner.fetch_plan(plan_ptr.id_at_location, alice)
+        alice = workers["alice"]
+        plan_ptr = plan.send(alice)
 
-    # Execute it locally
-    x = th.tensor([-1.26])
-    assert th.all(th.eq(fetched_plan(x), plan(x)))
-    # assert fetched_plan.state.state_ids != plan.state.state_ids #TODO
+        # Fetch plan
+        fetched_plan = plan.owner.fetch_plan(plan_ptr.id_at_location, alice)
 
-    # Make sure fetched_plan is using the readable_plan
-    assert fetched_plan.forward is None
-    assert fetched_plan.is_built
+        # Execute it locally
+        x = th.tensor([-1.26])
+        assert th.all(th.eq(fetched_plan(x), plan(x)))
+        # assert fetched_plan.state.state_ids != plan.state.state_ids #TODO
 
-    # Make sure plan is using the blueprint: forward
-    assert plan.forward is not None
+        # Make sure fetched_plan is using the readable_plan
+        assert fetched_plan.forward is None
+        assert fetched_plan.is_built
 
-    hook.local_worker.is_client_worker = True
+        # Make sure plan is using the blueprint: forward
+        assert plan.forward is not None
 
 
 @pytest.mark.parametrize("is_func2plan", [True, False])
 def test_fetch_stateful_plan_remote(hook, is_func2plan, start_remote_worker):
-    hook.local_worker.is_client_worker = False
 
-    server, remote_proxy = start_remote_worker(
-        id="test_fetch_stateful_plan_remote_{}".format(is_func2plan), hook=hook, port=8802
-    )
+    with hook.local_worker.registration_enabled():
+        server, remote_proxy = start_remote_worker(
+            id="test_fetch_stateful_plan_remote_{}".format(is_func2plan), hook=hook, port=8802
+        )
 
-    if is_func2plan:
+        if is_func2plan:
 
-        @sy.func2plan(args_shape=[(1,)], state=(th.tensor([3.0]),))
-        def plan(data, state):
-            bias, = state.read()
-            return data * bias
+            @sy.func2plan(args_shape=[(1,)], state=(th.tensor([3.0]),))
+            def plan(data, state):
+                bias, = state.read()
+                return data * bias
 
-    else:
+        else:
 
-        class Net(sy.Plan):
-            def __init__(self):
-                super(Net, self).__init__()
-                self.fc1 = nn.Linear(1, 1)
+            class Net(sy.Plan):
+                def __init__(self):
+                    super(Net, self).__init__()
+                    self.fc1 = nn.Linear(1, 1)
 
-            def forward(self, x):
-                return self.fc1(x)
-
-        plan = Net()
-        plan.build(th.tensor([1.2]))
+                def forward(self, x):
+                    return self.fc1(x)
 
-    x = th.tensor([-1.26])
-    expected = plan(x)
-    plan_ptr = plan.send(remote_proxy)
+            plan = Net()
+            plan.build(th.tensor([1.2]))
 
-    # Fetch plan
-    fetched_plan = plan.owner.fetch_plan(plan_ptr.id_at_location, remote_proxy)
+        x = th.tensor([-1.26])
+        expected = plan(x)
+        plan_ptr = plan.send(remote_proxy)
 
-    # Execute it locally
-    assert th.all(th.eq(fetched_plan(x), expected))
-    # assert fetched_plan.state.state_ids != plan.state.state_ids #TODO
+        # Fetch plan
+        fetched_plan = plan.owner.fetch_plan(plan_ptr.id_at_location, remote_proxy)
 
-    # Make sure fetched_plan is using the readable_plan
-    assert fetched_plan.forward is None
-    assert fetched_plan.is_built
+        # Execute it locally
+        assert th.all(th.eq(fetched_plan(x), expected))
+        # assert fetched_plan.state.state_ids != plan.state.state_ids #TODO
 
-    # Make sure plan is using the blueprint: forward
-    assert plan.forward is not None
+        # Make sure fetched_plan is using the readable_plan
+        assert fetched_plan.forward is None
+        assert fetched_plan.is_built
 
-    remote_proxy.close()
-    server.terminate()
+        # Make sure plan is using the blueprint: forward
+        assert plan.forward is not None
 
-    hook.local_worker.is_client_worker = True
+        remote_proxy.close()
+        server.terminate()
 
 
 def test_binding_fix_precision_plan(hook):
     """Here we make sure the attributes of a plan are still bound to state elements when calling fix_precision"""
 
-    hook.local_worker.is_client_worker = False
+    with hook.local_worker.registration_enabled():
 
-    class Net(sy.Plan):
-        def __init__(self):
-            super(Net, self).__init__()
-            self.fc1 = nn.Linear(1, 1)
-
-        def forward(self, x):
-            return self.fc1(x)
+        class Net(sy.Plan):
+            def __init__(self):
+                super(Net, self).__init__()
+                self.fc1 = nn.Linear(1, 1)
 
-    plan = Net()
-    plan.build(th.tensor([1.2]))
-    original_weight = plan.fc1.weight.clone()
+            def forward(self, x):
+                return self.fc1(x)
 
-    plan.fix_precision()
-    weight_id = plan.fc1.weight.id
-    hook.local_worker.get_obj(weight_id).float_prec_()
+        plan = Net()
+        plan.build(th.tensor([1.2]))
+        original_weight = plan.fc1.weight.clone()
 
-    assert (plan.fc1.weight - original_weight) < 10e-2
+        plan.fix_precision()
+        weight_id = plan.fc1.weight.id
+        hook.local_worker.get_obj(weight_id).float_prec_()
 
-    hook.local_worker.is_client_worker = True
+        assert (plan.fc1.weight - original_weight) < 10e-2
 
 
 def test_binding_encrypted_plan(hook, workers):
     """Here we make sure the attributes of a plan are still bound to state elements when calling fix_prec + share"""
 
-    hook.local_worker.is_client_worker = False
-
-    alice, bob, charlie, james = (
-        workers["alice"],
-        workers["bob"],
-        workers["charlie"],
-        workers["james"],
-    )
+    with hook.local_worker.registration_enabled():
 
-    class Net(sy.Plan):
-        def __init__(self):
-            super(Net, self).__init__()
-            self.fc1 = nn.Linear(1, 1)
+        alice, bob, charlie, james = (
+            workers["alice"],
+            workers["bob"],
+            workers["charlie"],
+            workers["james"],
+        )
 
-        def forward(self, x):
-            return self.fc1(x)
+        class Net(sy.Plan):
+            def __init__(self):
+                super(Net, self).__init__()
+                self.fc1 = nn.Linear(1, 1)
 
-    plan = Net()
-    plan.build(th.tensor([1.2]))
-    original_weight = plan.fc1.weight.clone()
+            def forward(self, x):
+                return self.fc1(x)
 
-    plan.fix_precision().share(alice, bob, crypto_provider=charlie)
-    weight_id = plan.fc1.weight.id
-    hook.local_worker.get_obj(weight_id).get_().float_prec_()
+        plan = Net()
+        plan.build(th.tensor([1.2]))
+        original_weight = plan.fc1.weight.clone()
 
-    assert (plan.fc1.weight - original_weight) < 10e-2
+        plan.fix_precision().share(alice, bob, crypto_provider=charlie)
+        weight_id = plan.fc1.weight.id
+        hook.local_worker.get_obj(weight_id).get_().float_prec_()
 
-    hook.local_worker.is_client_worker = True
+        assert (plan.fc1.weight - original_weight) < 10e-2
 
 
 @pytest.mark.parametrize("is_func2plan", [True, False])
 def test_fetch_encrypted_stateful_plan(hook, is_func2plan, workers):
     # TODO: this test is not working properly with remote workers.
     # We need to investigate why this might be the case.
-    hook.local_worker.is_client_worker = False
 
-    alice, bob, charlie, james = (
-        workers["alice"],
-        workers["bob"],
-        workers["charlie"],
-        workers["james"],
-    )
+    with hook.local_worker.registration_enabled():
+        alice, bob, charlie, james = (
+            workers["alice"],
+            workers["bob"],
+            workers["charlie"],
+            workers["james"],
+        )
 
-    if is_func2plan:
+        if is_func2plan:
 
-        @sy.func2plan(args_shape=[(1,)], state=(th.tensor([3.0]),))
-        def plan(data, state):
-            bias, = state.read()
-            return data * bias
+            @sy.func2plan(args_shape=[(1,)], state=(th.tensor([3.0]),))
+            def plan(data, state):
+                bias, = state.read()
+                return data * bias
 
-    else:
+        else:
 
-        class Net(sy.Plan):
-            def __init__(self):
-                super(Net, self).__init__()
-                self.fc1 = nn.Linear(1, 1)
+            class Net(sy.Plan):
+                def __init__(self):
+                    super(Net, self).__init__()
+                    self.fc1 = nn.Linear(1, 1)
 
-            def forward(self, x):
-                return self.fc1(x)
+                def forward(self, x):
+                    return self.fc1(x)
 
-        plan = Net()
-        plan.build(th.tensor([1.2]))
-
-    x = th.tensor([-1.0])
-    expected = plan(x)
+            plan = Net()
+            plan.build(th.tensor([1.2]))
 
-    plan.fix_precision().share(alice, bob, crypto_provider=charlie)
-    ptr_plan = plan.send(james)
+        x = th.tensor([-1.0])
+        expected = plan(x)
 
-    # Fetch plan
-    fetched_plan = plan.owner.fetch_plan(ptr_plan.id_at_location, james)
+        plan.fix_precision().share(alice, bob, crypto_provider=charlie)
+        ptr_plan = plan.send(james)
 
-    # Execute the fetch plan
-    x = th.tensor([-1.0])
-    x_sh = x.fix_precision().share(alice, bob, crypto_provider=charlie)
-    decrypted = fetched_plan(x_sh).get().float_prec()
+        # Fetch plan
+        fetched_plan = plan.owner.fetch_plan(ptr_plan.id_at_location, james)
 
-    # Compare with local plan
-    assert th.all(decrypted - expected.detach() < 1e-2)
-    # assert fetched_plan.state.state_ids != plan.state.state_ids #TODO
+        # Execute the fetch plan
+        x = th.tensor([-1.0])
+        x_sh = x.fix_precision().share(alice, bob, crypto_provider=charlie)
+        decrypted = fetched_plan(x_sh).get().float_prec()
 
-    # Make sure fetched_plan is using the readable_plan
-    assert fetched_plan.forward is None
-    assert fetched_plan.is_built
+        # Compare with local plan
+        assert th.all(decrypted - expected.detach() < 1e-2)
+        # assert fetched_plan.state.state_ids != plan.state.state_ids #TODO
 
-    # Make sure plan is using the blueprint: forward
-    assert plan.forward is not None
+        # Make sure fetched_plan is using the readable_plan
+        assert fetched_plan.forward is None
+        assert fetched_plan.is_built
 
-    hook.local_worker.is_client_worker = True
+        # Make sure plan is using the blueprint: forward
+        assert plan.forward is not None
 
 
 @pytest.mark.parametrize("is_func2plan", [True, False])
 def test_fecth_plan_multiple_times(hook, is_func2plan, workers):
-    hook.local_worker.is_client_worker = False
-
-    alice, bob, charlie, james = (
-        workers["alice"],
-        workers["bob"],
-        workers["charlie"],
-        workers["james"],
-    )
 
-    if is_func2plan:
+    with hook.local_worker.registration_enabled():
+        alice, bob, charlie, james = (
+            workers["alice"],
+            workers["bob"],
+            workers["charlie"],
+            workers["james"],
+        )
 
-        @sy.func2plan(args_shape=[(1,)], state=(th.tensor([3.0]),))
-        def plan(data, state):
-            bias, = state.read()
-            return data * bias
+        if is_func2plan:
 
-    else:
+            @sy.func2plan(args_shape=[(1,)], state=(th.tensor([3.0]),))
+            def plan(data, state):
+                bias, = state.read()
+                return data * bias
 
-        class Net(sy.Plan):
-            def __init__(self):
-                super(Net, self).__init__()
-                self.fc1 = nn.Linear(1, 1)
+        else:
 
-            def forward(self, x):
-                return self.fc1(x)
+            class Net(sy.Plan):
+                def __init__(self):
+                    super(Net, self).__init__()
+                    self.fc1 = nn.Linear(1, 1)
 
-        plan = Net()
-        plan.build(th.tensor([1.2]))
+                def forward(self, x):
+                    return self.fc1(x)
 
-    plan_pointer = plan.fix_precision().share(alice, bob, crypto_provider=charlie).send(james)
+            plan = Net()
+            plan.build(th.tensor([1.2]))
 
-    # Fetch plan
-    fetched_plan = plan_pointer.owner.fetch_plan(plan_pointer.id_at_location, james, copy=True)
+        plan_pointer = plan.fix_precision().share(alice, bob, crypto_provider=charlie).send(james)
 
-    # Execute the fetch plan
-    x = th.tensor([-1.0])
-    x_sh = x.fix_precision().share(alice, bob, crypto_provider=charlie)
-    decrypted1 = fetched_plan(x_sh).get().float_prec()
+        # Fetch plan
+        fetched_plan = plan_pointer.owner.fetch_plan(plan_pointer.id_at_location, james, copy=True)
 
-    # 2. Re-fetch Plan
-    fetched_plan = plan_pointer.owner.fetch_plan(plan_pointer.id_at_location, james, copy=True)
+        # Execute the fetch plan
+        x = th.tensor([-1.0])
+        x_sh = x.fix_precision().share(alice, bob, crypto_provider=charlie)
+        decrypted1 = fetched_plan(x_sh).get().float_prec()
 
-    # Execute the fetch plan
-    x = th.tensor([-1.0])
-    x_sh = x.fix_precision().share(alice, bob, crypto_provider=charlie)
-    decrypted2 = fetched_plan(x_sh).get().float_prec()
+        # 2. Re-fetch Plan
+        fetched_plan = plan_pointer.owner.fetch_plan(plan_pointer.id_at_location, james, copy=True)
 
-    assert th.all(decrypted1 - decrypted2 < 1e-2)
+        # Execute the fetch plan
+        x = th.tensor([-1.0])
+        x_sh = x.fix_precision().share(alice, bob, crypto_provider=charlie)
+        decrypted2 = fetched_plan(x_sh).get().float_prec()
 
-    hook.local_worker.is_client_worker = True
+        assert th.all(decrypted1 - decrypted2 < 1e-2)
 
 
 def test_fetch_plan_remote(hook, start_remote_worker):
-    hook.local_worker.is_client_worker = False
-
-    server, remote_proxy = start_remote_worker(id="test_fetch_plan_remote", hook=hook, port=8803)
-
-    @sy.func2plan(args_shape=[(1,)], state=(th.tensor([1.0]),))
-    def plan_mult_3(data, state):
-        bias, = state.read()
-        return data * 3 + bias
+    with hook.local_worker.registration_enabled():
+        server, remote_proxy = start_remote_worker(
+            id="test_fetch_plan_remote", hook=hook, port=8803
+        )
 
-    plan_mult_3.send(remote_proxy)
+        @sy.func2plan(args_shape=[(1,)], state=(th.tensor([1.0]),))
+        def plan_mult_3(data, state):
+            bias, = state.read()
+            return data * 3 + bias
 
-    # Fetch plan
-    fetched_plan = plan_mult_3.owner.fetch_plan(plan_mult_3.id, remote_proxy)
+        plan_mult_3.send(remote_proxy)
 
-    # Execute it locally
-    x = th.tensor([-1.0, 2, 3])
-    assert (plan_mult_3(x) == th.tensor([-2.0, 7, 10])).all()
-    assert (fetched_plan(x) == th.tensor([-2.0, 7, 10])).all()
-    assert fetched_plan.forward is None
-    assert fetched_plan.is_built
+        # Fetch plan
+        fetched_plan = plan_mult_3.owner.fetch_plan(plan_mult_3.id, remote_proxy)
 
-    remote_proxy.close()
-    server.terminate()
+        # Execute it locally
+        x = th.tensor([-1.0, 2, 3])
+        assert (plan_mult_3(x) == th.tensor([-2.0, 7, 10])).all()
+        assert (fetched_plan(x) == th.tensor([-2.0, 7, 10])).all()
+        assert fetched_plan.forward is None
+        assert fetched_plan.is_built
 
-    hook.local_worker.is_client_worker = True
+        remote_proxy.close()
+        server.terminate()
 
 
 def test_plan_serde(hook):
-    hook.local_worker.is_client_worker = False
-
-    @sy.func2plan(args_shape=[(1, 3)])
-    def my_plan(data):
-        x = data * 2
-        y = (x - 2) * 10
-        return x + y
+    with hook.local_worker.registration_enabled():
 
-    serialized_plan = serialize(my_plan)
-    deserialized_plan = deserialize(serialized_plan)
+        @sy.func2plan(args_shape=[(1, 3)])
+        def my_plan(data):
+            x = data * 2
+            y = (x - 2) * 10
+            return x + y
 
-    x = th.tensor([-1, 2, 3])
-    assert (deserialized_plan(x) == th.tensor([-42, 24, 46])).all()
+        serialized_plan = serialize(my_plan)
+        deserialized_plan = deserialize(serialized_plan)
 
-    hook.local_worker.is_client_worker = True
+        x = th.tensor([-1, 2, 3])
+        assert (deserialized_plan(x) == th.tensor([-42, 24, 46])).all()
 
 
 def test_execute_plan_remotely(hook, start_remote_worker):
@@ -711,114 +680,110 @@ def my_plan(data):
 def test_execute_plan_module_remotely(hook, start_remote_worker):
     """Test plan execution remotely."""
 
-    hook.local_worker.is_client_worker = False
-
-    class Net(sy.Plan):
-        def __init__(self):
-            super(Net, self).__init__()
-            self.fc1 = nn.Linear(2, 3)
-            self.fc2 = nn.Linear(3, 2)
+    with hook.local_worker.registration_enabled():
 
-            self.bias = th.tensor([1000.0])
+        class Net(sy.Plan):
+            def __init__(self):
+                super(Net, self).__init__()
+                self.fc1 = nn.Linear(2, 3)
+                self.fc2 = nn.Linear(3, 2)
 
-        def forward(self, x):
-            x = F.relu(self.fc1(x))
-            x = self.fc2(x)
-            return F.log_softmax(x, dim=0) + self.bias
+                self.bias = th.tensor([1000.0])
 
-    net = Net()
+            def forward(self, x):
+                x = F.relu(self.fc1(x))
+                x = self.fc2(x)
+                return F.log_softmax(x, dim=0) + self.bias
 
-    x = th.tensor([-1, 2.0])
-    local_res = net(x)
-    assert not net.is_built
+        net = Net()
 
-    net.build(x)
+        x = th.tensor([-1, 2.0])
+        local_res = net(x)
+        assert not net.is_built
 
-    server, remote_proxy = start_remote_worker(id="test_plan_worker_2", port=8799, hook=hook)
+        net.build(x)
 
-    plan_ptr = net.send(remote_proxy)
-    x_ptr = x.send(remote_proxy)
-    ptr = plan_ptr(x_ptr)
-    assert isinstance(ptr, FrameworkTensor) and ptr.is_wrapper
-    remote_res = ptr.get()
+        server, remote_proxy = start_remote_worker(id="test_plan_worker_2", port=8799, hook=hook)
 
-    assert (remote_res == local_res).all()
+        plan_ptr = net.send(remote_proxy)
+        x_ptr = x.send(remote_proxy)
+        ptr = plan_ptr(x_ptr)
+        assert isinstance(ptr, FrameworkTensor) and ptr.is_wrapper
+        remote_res = ptr.get()
 
-    # delete remote object before websocket connection termination
-    del x_ptr
+        assert (remote_res == local_res).all()
 
-    remote_proxy.close()
-    server.terminate()
+        # delete remote object before websocket connection termination
+        del x_ptr
 
-    hook.local_worker.is_client_worker = True
+        remote_proxy.close()
+        server.terminate()
 
 
 def test_train_plan_locally_and_then_send_it(hook, start_remote_worker):
     """Test training a plan locally and then executing it remotely."""
 
-    hook.local_worker.is_client_worker = False
-
-    # Create toy model
-    class Net(sy.Plan):
-        def __init__(self):
-            super(Net, self).__init__()
-            self.fc1 = nn.Linear(2, 3)
-            self.fc2 = nn.Linear(3, 2)
+    with hook.local_worker.registration_enabled():
 
-        def forward(self, x):
-            x = F.relu(self.fc1(x))
-            x = self.fc2(x)
-            return F.log_softmax(x, dim=0)
+        # Create toy model
+        class Net(sy.Plan):
+            def __init__(self):
+                super(Net, self).__init__()
+                self.fc1 = nn.Linear(2, 3)
+                self.fc2 = nn.Linear(3, 2)
 
-    net = Net()
+            def forward(self, x):
+                x = F.relu(self.fc1(x))
+                x = self.fc2(x)
+                return F.log_softmax(x, dim=0)
 
-    # Create toy data
-    x = th.tensor([-1, 2.0])
-    y = th.tensor([1.0])
+        net = Net()
 
-    # Train Model
-    opt = optim.SGD(params=net.parameters(), lr=0.01)
-    previous_loss = None
+        # Create toy data
+        x = th.tensor([-1, 2.0])
+        y = th.tensor([1.0])
 
-    for _ in range(5):
-        # 1) erase previous gradients (if they exist)
-        opt.zero_grad()
+        # Train Model
+        opt = optim.SGD(params=net.parameters(), lr=0.01)
+        previous_loss = None
 
-        # 2) make a prediction
-        pred = net(x)
+        for _ in range(5):
+            # 1) erase previous gradients (if they exist)
+            opt.zero_grad()
 
-        # 3) calculate how much we missed
-        loss = ((pred - y) ** 2).sum()
+            # 2) make a prediction
+            pred = net(x)
 
-        # 4) figure out which weights caused us to miss
-        loss.backward()
+            # 3) calculate how much we missed
+            loss = ((pred - y) ** 2).sum()
 
-        # 5) change those weights
-        opt.step()
+            # 4) figure out which weights caused us to miss
+            loss.backward()
 
-        if previous_loss is not None:
-            assert loss < previous_loss
+            # 5) change those weights
+            opt.step()
 
-        previous_loss = loss
+            if previous_loss is not None:
+                assert loss < previous_loss
 
-    local_res = net(x)
-    net.build(x)
+            previous_loss = loss
 
-    server, remote_proxy = start_remote_worker(id="test_plan_worker_3", port=8800, hook=hook)
+        local_res = net(x)
+        net.build(x)
 
-    plan_ptr = net.send(remote_proxy)
-    x_ptr = x.send(remote_proxy)
-    remote_res = plan_ptr(x_ptr).get()
+        server, remote_proxy = start_remote_worker(id="test_plan_worker_3", port=8800, hook=hook)
 
-    assert (remote_res == local_res).all()
+        plan_ptr = net.send(remote_proxy)
+        x_ptr = x.send(remote_proxy)
+        remote_res = plan_ptr(x_ptr).get()
 
-    # delete remote object before websocket connection termination
-    del x_ptr
+        assert (remote_res == local_res).all()
 
-    remote_proxy.close()
-    server.terminate()
+        # delete remote object before websocket connection termination
+        del x_ptr
 
-    hook.local_worker.is_client_worker = True
+        remote_proxy.close()
+        server.terminate()
 
 
 # def test_replace_worker_ids_two_strings(hook):
diff --git a/test/torch/tensors/test_precision.py b/test/torch/tensors/test_precision.py
--- a/test/torch/tensors/test_precision.py
+++ b/test/torch/tensors/test_precision.py
@@ -31,13 +31,11 @@ def test_encode_decode(workers, parameter):
 
 
 def test_fix_prec_registration(hook):
-    hook.local_worker.is_client_worker = False
+    with hook.local_worker.registration_enabled():
+        x = torch.tensor([1.0])
+        x_fpt = x.fix_precision()
 
-    x = torch.tensor([1.0])
-    x_fpt = x.fix_precision()
-    assert hook.local_worker.get_obj(x.id) == x
-
-    hook.local_worker.is_client_worker = True
+        assert hook.local_worker.get_obj(x.id) == x
 
 
 def test_inplace_encode_decode(workers):
@@ -54,13 +52,11 @@ def test_inplace_encode_decode(workers):
 
 
 def test_fix_prec_inplace_registration(hook):
-    hook.local_worker.is_client_worker = False
-
-    x = torch.tensor([1.0])
-    x.fix_precision_()
-    assert hook.local_worker.get_obj(x.id) == torch.tensor([1.0]).fix_precision()
 
-    hook.local_worker.is_client_worker = True
+    with hook.local_worker.registration_enabled():
+        x = torch.tensor([1.0])
+        x.fix_precision_()
+        assert hook.local_worker.get_obj(x.id) == torch.tensor([1.0]).fix_precision()
 
 
 def test_add_method():
diff --git a/test/workers/test_base.py b/test/workers/test_base.py
--- a/test/workers/test_base.py
+++ b/test/workers/test_base.py
@@ -88,3 +88,10 @@ def test_execute_command_self(hook):
     assert response == "bob_mocked_function"
 
     bob.mocked_function.assert_called()
+
+
+def test_enable_registration_with_ctx(hook):
+    assert hook.local_worker.is_client_worker == True
+    with hook.local_worker.registration_enabled():
+        hook.local_worker.is_client_worker == False
+    assert hook.local_worker.is_client_worker == True