Update RetrievalKit field tests to be more robust; fix duplicate decryption of treasure map for tests.

tests/unit/test_porter.py

@@ -100,20 +100,28 @@ def test_ursula_checksum_address_string_list_field(get_random_checksum_address):
         field._deserialize(value=f"{ursula_1},{ursula_2},{ursula_3},{ursula_4},0xdeadbeef", attr=None, data=None)
 
 
-def test_retrieval_kit_field():
-    encrypting_key = SecretKey.random().public_key()
-    capsule, _ = encrypt(encrypting_key, b'testing_retrieval_kit')
-
-    retrieval_kit = RetrievalKitClass(capsule, set())
-
+def test_retrieval_kit_field(get_random_checksum_address):
     field = RetrievalKit()
-    serialized = field._serialize(value=retrieval_kit, attr=None, obj=None)
-    assert serialized == b64encode(bytes(retrieval_kit)).decode()
 
-    deserialized = field._deserialize(value=serialized, attr=None, data=None)
-    assert isinstance(deserialized, RetrievalKitClass)
-    assert deserialized.capsule == capsule
-    assert len(deserialized.queried_addresses) == 0
+    def run_tests_on_kit(kit: RetrievalKitClass):
+        serialized = field._serialize(value=kit, attr=None, obj=None)
+        assert serialized == b64encode(bytes(kit)).decode()
+
+        deserialized = field._deserialize(value=serialized, attr=None, data=None)
+        assert isinstance(deserialized, RetrievalKitClass)
+        assert deserialized.capsule == kit.capsule
+        assert deserialized.queried_addresses == kit.queried_addresses
+
+    # kit with list of ursulas
+    encrypting_key = SecretKey.random().public_key()
+    capsule, _ = encrypt(encrypting_key, b'testing retrieval kit with 2 ursulas')
+    ursulas = [get_random_checksum_address(), get_random_checksum_address()]
+    run_tests_on_kit(kit=RetrievalKitClass(capsule, ursulas))
+
+    # kit with no ursulas
+    encrypting_key = SecretKey.random().public_key()
+    capsule, _ = encrypt(encrypting_key, b'testing retrieval kit with no ursulas')
+    run_tests_on_kit(kit=RetrievalKitClass(capsule, []))
 
     with pytest.raises(InvalidInputData):
         field._deserialize(value=b"non_base_64_data", attr=None, data=None)

tests/utils/policy.py

@@ -40,16 +40,12 @@ def generate_random_label() -> bytes:
 
 
 def retrieval_request_setup(enacted_policy, bob, alice, encode_for_rest=False) -> Dict:
-
     treasure_map = bob._decrypt_treasure_map(enacted_policy.treasure_map)
 
     # We pick up our story with Bob already having followed the treasure map above, ie:
     bob.start_learning_loop()
 
     bob.follow_treasure_map(treasure_map=treasure_map, block=True, timeout=1)
-    decrypted_treasure_map = bob._decrypt_treasure_map(enacted_policy.treasure_map,
-                                                       enacted_policy.publisher_verifying_key)
-
     # We'll test against just a single Ursula - here, we make a WorkOrder for just one.
     # We can pass any number of capsules as args; here we pass just one.
     enrico = Enrico(policy_encrypting_key=enacted_policy.public_key)
@@ -58,7 +54,7 @@ def retrieval_request_setup(enacted_policy, bob, alice, encode_for_rest=False) -
 
     encode_bytes = (lambda field, obj: field()._serialize(value=obj, attr=None, obj=None)) if encode_for_rest else (lambda field, obj: obj)
 
-    return dict(treasure_map=encode_bytes(TreasureMap, decrypted_treasure_map),
+    return dict(treasure_map=encode_bytes(TreasureMap, treasure_map),
                 retrieval_kits=[encode_bytes(RetrievalKit, message_kit.as_retrieval_kit())],
                 alice_verifying_key=encode_bytes(Key, alice.stamp.as_umbral_pubkey()),
                 bob_encrypting_key=encode_bytes(Key, bob.public_keys(DecryptingPower)),