More flexible address book

core/capabilities/ccip/deployment/deploy.go

@@ -57,13 +57,20 @@ func GenerateJobSpecs(capReg common.Address) CCIPSpec {
 	return CCIPSpec{}
 }
 
+type DeployCCIPContractConfig struct {
+	Weth9s map[uint64]common.Address
+	// TODO: More params as needed
+}
+
 // TODO: Likely we'll want to further parameterize the deployment
 // For example a list of contracts to skip deploying if they already exist.
 // Or mock vs real RMN.
-func DeployCCIPContracts(e environment.Environment) error {
+// Deployment produces an address book of everything it deployed.
+func DeployCCIPContracts(e environment.Environment, c DeployCCIPContractConfig) (environment.AddressBook, error) {
+	ab := environment.NewMemoryAddressBook()
 	for _, chain := range e.Chains {
 		saveToChain := func(addr common.Address) error {
-			return e.AddressBook.Save(chain.Selector, addr.String())
+			return ab.Save(chain.Selector, addr.String())
 		}
 
 		// TODO: Still waiting for RMNRemote/RMNHome contracts etc.
@@ -77,7 +84,7 @@ func DeployCCIPContracts(e environment.Environment) error {
 			}, chain.Confirm, saveToChain)
 		if err != nil {
 			e.Logger.Errorw("Failed to deploy mockARM", "err", err)
-			return err
+			return ab, err
 		}
 		e.Logger.Infow("deployed mockARM", "addr", mockARM)
 
@@ -92,7 +99,7 @@ func DeployCCIPContracts(e environment.Environment) error {
 			}, chain.Confirm, saveToChain)
 		if err != nil {
 			e.Logger.Errorw("Failed to deploy armProxy", "err", err)
-			return err
+			return ab, err
 		}
 		e.Logger.Infow("deployed armProxy", "addr", armProxy)
 
@@ -121,7 +128,7 @@ func DeployCCIPContracts(e environment.Environment) error {
 			}, chain.Confirm, saveToChain)
 		if err != nil {
 			e.Logger.Errorw("Failed to deploy router", "err", err)
-			return err
+			return ab, err
 		}
 		e.Logger.Infow("deployed router", "addr", routerAddr)
 
@@ -134,9 +141,9 @@ func DeployCCIPContracts(e environment.Environment) error {
 			}, chain.Confirm, saveToChain)
 		if err != nil {
 			e.Logger.Errorw("Failed to deploy token admin registry", "err", err)
-			return err
+			return ab, err
 		}
 		e.Logger.Infow("deployed tokenAdminRegistry", "addr", tokenAdminRegistry)
 	}
-	return nil
+	return ab, nil
 }

core/capabilities/ccip/deployment/deploy_test.go

@@ -16,15 +16,16 @@ func TestDeployCCIPContracts(t *testing.T) {
 		Nodes:  1,
 	})
 	// Deploy all the CCIP contracts.
-	require.NoError(t, DeployCCIPContracts(e))
-	state, err := GenerateOnchainState(e)
+	ab, err := DeployCCIPContracts(e, DeployCCIPContractConfig{})
+	require.NoError(t, err)
+	state, err := GenerateOnchainState(e, ab)
 	require.NoError(t, err)
 	snap, err := state.Snapshot(e.AllChainSelectors())
 	require.NoError(t, err)
 
 	// Assert expect every deployed address to be in the address book.
 	for name, chain := range snap.Chains {
-		addrs, err := e.AddressBook.Addresses()
+		addrs, err := ab.Addresses()
 		require.NoError(t, err)
 		evmChainID, _ := chainsel.ChainIdFromName(name)
 		sel, _ := chainsel.SelectorFromChainId(evmChainID)

core/capabilities/ccip/deployment/state.go

@@ -70,7 +70,15 @@ func (s CCIPOnChainState) Snapshot(chains []uint64) (CCIPSnapShot, error) {
 	return snapshot, nil
 }
 
-func GenerateOnchainState(e environment.Environment) (CCIPOnChainState, error) {
+func SnapshotState(e environment.Environment, ab environment.AddressBook) (CCIPSnapShot, error) {
+	state, err := GenerateOnchainState(e, ab)
+	if err != nil {
+		return CCIPSnapShot{}, err
+	}
+	return state.Snapshot(e.AllChainSelectors())
+}
+
+func GenerateOnchainState(e environment.Environment, ab environment.AddressBook) (CCIPOnChainState, error) {
 	state := CCIPOnChainState{
 		EvmOnRampsV160:       make(map[uint64]*evm_2_evm_multi_onramp.EVM2EVMMultiOnRamp),
 		EvmOffRampsV160:      make(map[uint64]*evm_2_evm_multi_offramp.EVM2EVMMultiOffRamp),
@@ -82,7 +90,7 @@ func GenerateOnchainState(e environment.Environment) (CCIPOnChainState, error) {
 		Weth9s:               make(map[uint64]*weth9.WETH9),
 	}
 	// Get all the onchain state
-	addresses, err := e.AddressBook.Addresses()
+	addresses, err := ab.Addresses()
 	if err != nil {
 		return state, errors.Wrap(err, "could not get addresses")
 	}

core/environment/address_book.go

@@ -0,0 +1,47 @@
+package environment
+
+type AddressBook interface {
+	// TODO: Need manualTV override
+	Save(chainSelector uint64, address string) error
+	Addresses() (map[uint64]map[string]struct{}, error)
+	// Allows for merging address books
+	Merge(other AddressBook) error
+}
+
+type AddressBookMap struct {
+	AddressesByChain map[uint64]map[string]struct{}
+}
+
+func (m *AddressBookMap) Save(chainSelector uint64, address string) error {
+	if _, exists := m.AddressesByChain[chainSelector]; !exists {
+		m.AddressesByChain[chainSelector] = make(map[string]struct{})
+	} else {
+		// Error?
+	}
+	m.AddressesByChain[chainSelector][address] = struct{}{}
+	return nil
+}
+
+func (m *AddressBookMap) Addresses() (map[uint64]map[string]struct{}, error) {
+	return m.AddressesByChain, nil
+}
+
+// Attention this will mutate existing book
+func (m *AddressBookMap) Merge(ab AddressBook) error {
+	addresses, err := ab.Addresses()
+	if err != nil {
+		return err
+	}
+	for chain, chainAddresses := range addresses {
+		for address := range chainAddresses {
+			return m.Save(chain, address)
+		}
+	}
+	return nil
+}
+
+func NewMemoryAddressBook() *AddressBookMap {
+	return &AddressBookMap{
+		AddressesByChain: make(map[uint64]map[string]struct{}),
+	}
+}

core/environment/address_book_test.go

@@ -0,0 +1 @@
+package environment

core/environment/environment.go

@@ -15,11 +15,9 @@ import (
 	"github.com/smartcontractkit/chainlink/v2/core/environment/memory"
 )
 
-type AddressBook interface {
-	// TODO: Need manualTV override
-	Save(chainSelector uint64, address string) error
-	Addresses() (map[uint64]map[string]struct{}, error)
-}
+const (
+	Memory = "memory"
+)
 
 type OnchainClient interface {
 	// For EVM specifically we can use existing geth interface
@@ -43,11 +41,11 @@ type Chain struct {
 }
 
 type Environment struct {
-	Chains      map[uint64]Chain
-	NodeIds     []string
-	AddressBook AddressBook
-	Offchain    OffchainClient
-	Logger      logger.Logger
+	Name     string
+	Chains   map[uint64]Chain
+	NodeIds  []string
+	Offchain OffchainClient
+	Logger   logger.Logger
 }
 
 func (e Environment) AllChainSelectors() []uint64 {
@@ -92,10 +90,10 @@ func NewMemoryEnvironment(t *testing.T, config MemoryEnvironmentConfig) Environm
 	lggr, err := logger.New()
 	require.NoError(t, err)
 	return Environment{
-		Offchain:    memory.NewMemoryJobClient(nodesByPeerID),
-		Chains:      chains,
-		NodeIds:     keys,
-		AddressBook: memory.NewMemoryAddressBook(),
-		Logger:      lggr,
+		Name:     Memory,
+		Offchain: memory.NewMemoryJobClient(nodesByPeerID),
+		Chains:   chains,
+		NodeIds:  keys,
+		Logger:   lggr,
 	}
 }