diff options
Diffstat (limited to 'swarm/network')
| -rw-r--r-- | swarm/network/simulations/discovery/discovery_test.go | 31 |
1 files changed, 25 insertions, 6 deletions
diff --git a/swarm/network/simulations/discovery/discovery_test.go b/swarm/network/simulations/discovery/discovery_test.go index 3c3affe58..cd5456b73 100644 --- a/swarm/network/simulations/discovery/discovery_test.go +++ b/swarm/network/simulations/discovery/discovery_test.go @@ -85,11 +85,12 @@ func getDbStore(nodeID string) (*state.DBStore, error) { } var ( - nodeCount = flag.Int("nodes", 10, "number of nodes to create (default 10)") - initCount = flag.Int("conns", 1, "number of originally connected peers (default 1)") - snapshotFile = flag.String("snapshot", "", "create snapshot") - loglevel = flag.Int("loglevel", 3, "verbosity of logs") - rawlog = flag.Bool("rawlog", false, "remove terminal formatting from logs") + nodeCount = flag.Int("nodes", 10, "number of nodes to create (default 10)") + initCount = flag.Int("conns", 1, "number of originally connected peers (default 1)") + snapshotFile = flag.String("snapshot", "", "path to create snapshot file in") + loglevel = flag.Int("loglevel", 3, "verbosity of logs") + rawlog = flag.Bool("rawlog", false, "remove terminal formatting from logs") + serviceOverride = flag.String("services", "", "remove or add services to the node snapshot; prefix with \"+\" to add, \"-\" to remove; example: +pss,-discovery") ) func init() { @@ -306,7 +307,25 @@ func discoverySimulation(nodes, conns int, adapter adapters.NodeAdapter) (*simul } if *snapshotFile != "" { - snap, err := net.Snapshot() + var err error + var snap *simulations.Snapshot + if len(*serviceOverride) > 0 { + var addServices []string + var removeServices []string + for _, osvc := range strings.Split(*serviceOverride, ",") { + if strings.Index(osvc, "+") == 0 { + addServices = append(addServices, osvc[1:]) + } else if strings.Index(osvc, "-") == 0 { + removeServices = append(removeServices, osvc[1:]) + } else { + panic("stick to the rules, you know what they are") + } + } + snap, err = net.SnapshotWithServices(addServices, removeServices) + } else { + snap, err = net.Snapshot() + } + if err != nil { return nil, errors.New("no shapshot dude") } |