Have lookup_contract search better

nucypher/blockchain/eth/interfaces.py

@@ -132,8 +132,16 @@ class EthereumContractRegistrar:
             raise self.UnknownChain("Data does not exist for chain '{}'".format(self._chain_name))
         return chain_data
 
-    def lookup_contract(self, contract_name: str):
+    def lookup_contract(self, contract_name: str=None, contract_addr: str=None):
+        if not (bool(contract_name) ^ bool(contract_addr)):
+            raise ValueError("Pass contract_name or contract_addr, not both.")
+
         chain_data = self.dump_chain()
+        if contract_addr:
+            contract_data = chain_data.get(contract_addr, None)
+            if not contract_data:
+                raise self.UnknownContract("No known contract with addr: {}".format(contract_addr))
+            return chain_data[contract_addr]
 
         for address, contract_data in chain_data.items():
             name = contract_data.get('name', None)

tests/blockchain/eth/client/test_interfaces.py

@@ -3,24 +3,6 @@ import pytest
 from nucypher.blockchain.eth.interfaces import EthereumContractRegistrar
 
 
-#def test_registrar_read_write(tempfile_path):
-#    # Test that the file is initally empty and returns an empty dict.
-#    should_be_empty = _read_registrar_file(tempfile_path)
-#    assert should_be_empty == {}
-#
-#    # Test that data can be written and read
-#    test_data = {'test': 'foobar'}
-#    _write_registrar_file(test_data, tempfile_path)
-#    out_data = _read_registrar_file(tempfile_path)
-#    assert test_data == out_data
-#
-#    # Test overwrite
-#    new_test_data = {'new_test': 'foobar'}
-#    _write_registrar_file(new_test_data, tempfile_path)
-#    out_data = _read_registrar_file(tempfile_path)
-#    assert out_data != test_data and new_test_data == out_data
-
-
 def test_registrar_object(tempfile_path):
     # Tests everything is as it should be when initially created
     test_registrar = EthereumContractRegistrar(registrar_filepath=tempfile_path)
@@ -33,14 +15,12 @@ def test_registrar_object(tempfile_path):
     assert isinstance(contains_registrar['tester'], EthereumContractRegistrar)
 
     # Test contract enrollment and dump_chain
-    test_name = 'test_contract'
     test_addr = '0xDEADBEEF'
     test_abi = ['fake', 'data']
-    test_registrar.enroll(test_name, test_addr, test_abi)
+    test_registrar.enroll(test_addr, test_abi)
 
     chain_data = test_registrar.dump_chain()
     assert test_addr in chain_data
-    assert chain_data[test_addr]['name'] == test_name
     assert chain_data[test_addr]['addr'] == test_addr
     assert chain_data[test_addr]['abi'] == test_abi
 
@@ -55,13 +35,22 @@ def test_registrar_object(tempfile_path):
 
     assert contract_by_name == contract_by_addr
 
-    # Test enroll with same contract name updates contract data
-    new_test_addr = '0xBEEFDEAD'
-    test_registrar.enroll(test_name, new_test_addr, test_abi)
+    # Test enrolling a dispatcher
+    new_test_addr = '0xTESTDISPATCHER'
+    new_test_abi = ['dispatch', 'test', 'info']
+    new_test_name = 'TestNameDispatcher'
+    new_test_target_addr = '0xTARGET'
+    test_registrar.enroll(new_test_addr, new_test_abi, new_test_name, new_test_target_addr)
 
-    test_contract = test_registrar.dump_contract(new_test_addr)
-    assert test_contract['addr'] == new_test_addr
-    assert test_contract['abi'] == test_abi
+    contract_by_name = test_registrar.lookup_contract(new_test_name)
+    assert contract_by_name['addr'] == new_test_addr
+    assert contract_by_name['abi'] == new_test_abi
+    assert contract_by_name['name'] == new_test_name
+    assert contract_by_name['target_addr'] == new_test_target_addr
+
+    # Check that it raises an error
+    with pytest.raises(EthereumContractRegistrar.UnknownContract):
+        test_registrar.lookup_contract('this will not exist')
 
     # Test new chain via new registrar object
     new_chain_name = 'not_tester'
@@ -69,7 +58,7 @@ def test_registrar_object(tempfile_path):
     chains = EthereumContractRegistrar.get_registrars(tempfile_path)
     assert new_chain_name not in chains # Not written yet, shouldn't be there
 
-    new_chain_registrar.enroll(new_chain_name, new_test_addr, test_abi)
+    new_chain_registrar.enroll(new_test_addr, test_abi)
     updated_chains = EthereumContractRegistrar.get_registrars(tempfile_path)
     assert new_chain_name in updated_chains and 'tester' in updated_chains