Saving replete WorkOrders where appropriate.

nucypher/characters/lawful.py

@@ -629,15 +629,14 @@ class Bob(Character):
 
         return incomplete_work_orders, complete_work_orders
 
-    def get_reencrypted_cfrags(self, work_order):
+    def get_reencrypted_cfrags(self, work_order, retain_cfrags=False):
         if work_order.completed:
             raise TypeError("This WorkOrder is already complete; if you want Ursula to perform additional service, make a new WorkOrder.")
 
-        cfrags = self.network_middleware.reencrypt(work_order)
+        cfrags_and_signatures = self.network_middleware.reencrypt(work_order)
+        cfrags = work_order.complete(cfrags_and_signatures)
+        self._completed_work_orders.save_work_order(work_order, as_replete=retain_cfrags)
 
-        for task in work_order.tasks.values():
-            completed_work_orders_for_ursula = self._completed_work_orders.by_checksum_address(work_order.ursula.checksum_address)
-            completed_work_orders_for_ursula[task.capsule] = work_order
         return cfrags
 
     def join_policy(self, label, alice_verifying_key, node_list=None, block=False):
@@ -689,13 +688,13 @@ class Bob(Character):
             # TODO Optimization: Block here (or maybe even later) until map is done being followed (instead of blocking above). #1114
             the_airing_of_grievances = []
 
-            for work_order in incomplete_work_orders.values():
+            for work_order in new_work_orders.values():
                 if len(capsule) >= m:
                     # TODO: What to do with unused WorkOrders here?   #1197
                     break
                 # We don't have enough CFrags yet.  Let's get another one from a WorkOrder.
                 try:
-                    cfrags = self.get_reencrypted_cfrags(work_order)
+                    cfrags = self.get_reencrypted_cfrags(work_order, retain_cfrags=retain_cfrags)
                 except NodeSeemsToBeDown:
                     # TODO: What to do here?  Ursula isn't supposed to be down.
                     self.log.info(

nucypher/network/middleware.py

@@ -198,8 +198,7 @@ class RestMiddleware:
         ursula_rest_response = self.send_work_order_payload_to_ursula(work_order)
         splitter = BytestringSplitter((CapsuleFrag, VariableLengthBytestring), Signature)
         cfrags_and_signatures = splitter.repeat(ursula_rest_response.content)
-        cfrags = work_order.complete(cfrags_and_signatures)
-        return cfrags
+        return cfrags_and_signatures
 
     def revoke_arrangement(self, ursula, revocation):
         # TODO: Implement revocation confirmations

nucypher/policy/collections.py

@@ -394,6 +394,7 @@ class WorkOrderHistory:
 
     def __init__(self) -> None:
         self.by_ursula = {}  # type: dict
+        self._latest_replete = {}
 
     def __contains__(self, item):
         assert False

tests/characters/test_bob_handles_frags.py

@@ -167,7 +167,7 @@ def test_bob_can_issue_a_work_order_to_a_specific_ursula(enacted_federated_polic
     assert work_order.completed is False
 
     # **** RE-ENCRYPTION HAPPENS HERE! ****
-    cfrags = federated_bob.get_reencrypted_cfrags(work_order)
+    cfrags = federated_bob.get_reencrypted_cfrags(work_order, retain_cfrags=True)
 
     # We only gave one Capsule, so we only got one cFrag.
     assert len(cfrags) == 1
@@ -224,8 +224,9 @@ def test_bob_can_use_cfrag_attached_to_completed_workorder(enacted_federated_pol
         num_ursulas=1,
         )
 
-    # Here we show that since we're using the same completed WorkOrder again, we get it back.
-    new_work_order = list(complete_work_orders.values())[0]
+    # Here we show that this WorkOrder is still saved, replete with the CFrag.
+    work_orders_for_this_capsule = federated_bob._completed_work_orders._latest_replete[last_capsule_on_side_channel]
+    new_work_order = list(work_orders_for_this_capsule.values())[0]
     assert old_work_order == new_work_order
 
     # We already got a CFrag for this WorkOrder, a couple of tests ago.
@@ -280,7 +281,7 @@ def test_bob_remembers_that_he_has_cfrags_for_a_particular_capsule(enacted_feder
     assert new_work_order != saved_work_order
 
     # This WorkOrder has never been completed
-    cfrags = federated_bob.get_reencrypted_cfrags(new_work_order)
+    cfrags = federated_bob.get_reencrypted_cfrags(new_work_order, retain_cfrags=True)
 
     # Again: one Capsule, one cFrag.
     assert len(cfrags) == 1