cmd/swarm/swarm-snapshot: refactor and extend tests

cmd/swarm/swarm-snapshot/create.go

@@ -23,14 +23,14 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
+	"path"
+	"path/filepath"
 	"strings"
 	"sync"
 	"time"
 
-	"github.com/ethereum/go-ethereum/cmd/utils"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/node"
-	"github.com/ethereum/go-ethereum/p2p/enode"
 	"github.com/ethereum/go-ethereum/p2p/simulations"
 	"github.com/ethereum/go-ethereum/p2p/simulations/adapters"
 	"github.com/ethereum/go-ethereum/swarm/network"
@@ -38,12 +38,11 @@ import (
 	cli "gopkg.in/urfave/cli.v1"
 )
 
-const testMinProxBinSize = 2
-const NoConnectionTimeout = 2 * time.Second
+const noConnectionTimeout = 2 * time.Second
 
 func create(ctx *cli.Context) error {
 	log.PrintOrigins(true)
-	log.Root().SetHandler(log.LvlFilterHandler(log.Lvl(verbosity), log.StreamHandler(os.Stdout, log.TerminalFormat(true))))
+	log.Root().SetHandler(log.LvlFilterHandler(log.Lvl(ctx.Int("verbosity")), log.StreamHandler(os.Stdout, log.TerminalFormat(true))))
 
 	if len(ctx.Args()) < 1 {
 		return errors.New("argument should be the filename to verify or write-to")
@@ -52,38 +51,19 @@ func create(ctx *cli.Context) error {
 	if err != nil {
 		return err
 	}
-	err = discoverySnapshot(filename, 10)
-	if err != nil {
-		utils.Fatalf("Simulation failed: %s", err)
+	return createSnapshot(filename, ctx.Int("nodes"), ctx.String("services"))
 }
 
-	return err
-}
+func createSnapshot(filename string, nodes int, services string) (err error) {
+	log.Debug("create snapshot", "filename", filename, "nodes", nodes, "services", services)
 
-func discoverySnapshot(filename string, nodes int) error {
-	//disable discovery if topology is specified
-	discovery = topology == ""
-	log.Debug("discoverySnapshot", "filename", filename, "nodes", nodes, "discovery", discovery)
-	i := 0
-	var lock sync.Mutex
-	var pivotNodeID enode.ID
 	sim := simulation.New(map[string]simulation.ServiceFunc{
-		"bzz": func(ctx *adapters.ServiceContext, b *sync.Map) (node.Service, func(), error) {
-			lock.Lock()
-			i++
-			if i == pivot {
-				pivotNodeID = ctx.Config.ID
-			}
-			lock.Unlock()
-
+		bzzServiceName: func(ctx *adapters.ServiceContext, b *sync.Map) (node.Service, func(), error) {
 			addr := network.NewAddr(ctx.Config.Node())
-			kp := network.NewKadParams()
-			kp.MinProxBinSize = testMinProxBinSize
-
-			kad := network.NewKademlia(addr.Over(), kp)
+			kad := network.NewKademlia(addr.Over(), network.NewKadParams())
 			hp := network.NewHiveParams()
 			hp.KeepAliveInterval = time.Duration(200) * time.Millisecond
-			hp.Discovery = discovery
+			hp.Discovery = true // discovery must be enabled when creating a snapshot
 
 			config := &network.BzzConfig{
 				OverlayAddr:  addr.Over(),
@@ -95,9 +75,9 @@ func discoverySnapshot(filename string, nodes int) error {
 	})
 	defer sim.Close()
 
-	_, err := sim.AddNodes(10)
+	_, err = sim.AddNodes(nodes)
 	if err != nil {
-		utils.Fatalf("%v", err)
+		return fmt.Errorf("add nodes: %v", err)
 	}
 
 	events := make(chan *simulations.Event)
@@ -106,60 +86,84 @@ func discoverySnapshot(filename string, nodes int) error {
 	case ev := <-events:
 		//only catch node up events
 		if ev.Type == simulations.EventTypeConn {
-			utils.Fatalf("this shouldn't happen as connections weren't initiated yet")
+			return errors.New("unexpected connection events")
 		}
-	case <-time.After(NoConnectionTimeout):
+	case <-time.After(noConnectionTimeout):
 	}
-
 	sub.Unsubscribe()
 
 	if len(sim.Net.Conns) > 0 {
-		utils.Fatalf("no connections should exist after just adding nodes")
+		return errors.New("no connections should exist after just adding nodes")
 	}
 
-	err := sim.Net.ConnectNodesRing(nil)
+	err = sim.Net.ConnectNodesRing(nil)
 	if err != nil {
-		utils.Fatalf("had an error connecting the nodes in a %v topology: %v", topology, err)
+		return fmt.Errorf("connect nodes: %v", err)
 	}
 
-	if discovery {
 	ctx, cancelSimRun := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancelSimRun()
-
 	if _, err := sim.WaitTillHealthy(ctx, 2); err != nil {
-			utils.Fatalf("%v", err)
-		}
+		return fmt.Errorf("wait for healthy kademlia: %v", err)
 	}
 
 	var snap *simulations.Snapshot
 	if len(services) > 0 {
-		var addServices []string
+		addServices := strings.Split(services, ",")
 		var removeServices []string
-		for _, osvc := range strings.Split(services, ",") {
-			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")
+		var hasBzz bool
+		for _, s := range addServices {
+			if s == bzzServiceName {
+				hasBzz = true
+				break
 			}
 		}
+		if !hasBzz {
+			removeServices = append(removeServices, bzzServiceName)
+		}
 		snap, err = sim.Net.SnapshotWithServices(addServices, removeServices)
 	} else {
 		snap, err = sim.Net.Snapshot()
 	}
-
 	if err != nil {
-		return errors.New("no shapshot dude")
+		return fmt.Errorf("create snapshot: %v", err)
 	}
 	jsonsnapshot, err := json.Marshal(snap)
 	if err != nil {
-		return fmt.Errorf("corrupt json snapshot: %v", err)
+		return fmt.Errorf("json encode snapshot: %v", err)
 	}
-	err = ioutil.WriteFile(filename, jsonsnapshot, 0666)
-	if err != nil {
-		return err
+	return ioutil.WriteFile(filename, jsonsnapshot, 0666)
 }
 
-	return nil
+func touchPath(filename string) (string, error) {
+	if path.IsAbs(filename) {
+		if _, err := os.Stat(filename); err == nil {
+			// path exists, we will override the file
+			return filename, nil
+		}
+	}
+
+	d, f := path.Split(filename)
+	dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
+	if err != nil {
+		return "", err
+	}
+
+	_, err = os.Stat(path.Join(dir, filename))
+	if err == nil {
+		// path exists, we will override
+		return filename, nil
+	}
+
+	dirPath := path.Join(dir, d)
+	filePath := path.Join(dirPath, f)
+	if d != "" {
+		err = os.MkdirAll(dirPath, os.ModeDir)
+		if err != nil {
+			return "", err
+		}
+	}
+
+	filename = filePath
+	return filename, nil
 }

cmd/swarm/swarm-snapshot/create_test.go

@@ -17,58 +17,121 @@
 package main
 
 import (
+	"encoding/json"
+	"fmt"
 	"io/ioutil"
 	"os"
+	"sort"
+	"strconv"
+	"strings"
 	"testing"
 
-	"github.com/ethereum/go-ethereum/log"
+	"github.com/ethereum/go-ethereum/p2p/simulations"
 )
 
-func init() {
-	log.PrintOrigins(true)
-	log.Root().SetHandler(log.LvlFilterHandler(log.Lvl(verbosity), log.StreamHandler(os.Stdout, log.TerminalFormat(true))))
-}
-
 //TestSnapshotCreate is a high level e2e test that tests for snapshot generation
 func TestSnapshotCreate(t *testing.T) {
-
 	for _, v := range []struct {
 		name     string
-		args []string
+		nodes    int
+		services string
 	}{
 		{
-			name: "no topology - discovery enabled",
-			args: []string{
-				"c",
-			},
+			name: "defaults",
 		},
 		{
-			name: "yes topology - discovery disabled",
-			args: []string{
-				"--topology",
-				"ring",
-				"c",
+			name:  "more nodes",
+			nodes: defaultNodes + 5,
 		},
+		{
+			name:     "services",
+			services: "stream,pss,zorglub",
+		},
+		{
+			name:     "services with " + bzzServiceName,
+			services: bzzServiceName + ",pss",
 		},
 	} {
 		t.Run(v.name, func(t *testing.T) {
+			t.Parallel()
+
 			file, err := ioutil.TempFile("", "swarm-snapshot")
+			if err != nil {
+				t.Fatal(err)
+			}
 			defer os.Remove(file.Name())
+
+			if err = file.Close(); err != nil {
+				t.Error(err)
+			}
+
+			args := []string{"create"}
+			if v.nodes > 0 {
+				args = append(args, "--nodes", strconv.Itoa(v.nodes))
+			}
+			if v.services != "" {
+				args = append(args, "--services", v.services)
+			}
+			testCmd := runSnapshot(t, append(args, file.Name())...)
+
+			testCmd.ExpectExit()
+			if testCmd.ExitStatus() != 0 {
+				t.Fatal("expected exit code 0")
+			}
+
+			f, err := os.Open(file.Name())
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer func() {
+				err := f.Close()
+				if err != nil {
+					t.Error("closing snapshot file", "err", err)
+				}
+			}()
+
+			b, err := ioutil.ReadAll(f)
+			if err != nil {
+				t.Fatal(err)
+			}
+			var snap simulations.Snapshot
+			err = json.Unmarshal(b, &snap)
 			if err != nil {
 				t.Fatal(err)
 			}
 
-			file.Close()
-			snap := runSnapshot(t, append(v.args, file.Name())...)
-
-			snap.ExpectExit()
-			if snap.ExitStatus() != 0 {
-				t.Fatal("expected exit code 0")
+			wantNodes := v.nodes
+			if wantNodes == 0 {
+				wantNodes = defaultNodes
+			}
+			gotNodes := len(snap.Nodes)
+			if gotNodes != wantNodes {
+				t.Errorf("got %v nodes, want %v", gotNodes, wantNodes)
 			}
 
-			_, err = os.Stat(file.Name())
+			if len(snap.Conns) == 0 {
+				t.Error("no connections in a snapshot")
+			}
+
+			var wantServices []string
+			if v.services != "" {
+				wantServices = strings.Split(v.services, ",")
+			} else {
+				wantServices = []string{bzzServiceName}
+			}
+			sort.Strings(wantServices)
+
+			for i, n := range snap.Nodes {
+				gotServices := n.Node.Config.Services
+				sort.Strings(gotServices)
+				if fmt.Sprint(gotServices) != fmt.Sprint(wantServices) {
+					t.Errorf("got services %v for node %v, want %v", gotServices, i, wantServices)
+				}
+			}
+
+			err = verifySnapshot(file.Name())
 			if err != nil {
-				t.Fatal("could not stat snapshot json")
+				t.Error(err)
 			}
 		})
 	}

cmd/swarm/swarm-snapshot/helper.go

@@ -1,40 +0,0 @@
-package main
-
-import (
-	"os"
-	"path"
-	"path/filepath"
-)
-
-func touchPath(filename string) (string, error) {
-	if path.IsAbs(filename) {
-		if _, err := os.Stat(filename); err == nil {
-			// path exists, we will override the file
-			return filename, nil
-		}
-	}
-
-	d, f := path.Split(filename)
-	dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
-	if err != nil {
-		return "", err
-	}
-
-	_, err = os.Stat(path.Join(dir, filename))
-	if err == nil {
-		// path exists, we will override
-		return filename, nil
-	}
-
-	dirPath := path.Join(dir, d)
-	filePath := path.Join(dirPath, f)
-	if d != "" {
-		err = os.MkdirAll(dirPath, os.ModeDir)
-		if err != nil {
-			return "", err
-		}
-	}
-
-	filename = filePath
-	return filename, nil
-}

cmd/swarm/swarm-snapshot/main.go

@@ -18,65 +18,38 @@ package main
 
 import (
 	"os"
-	"sort"
 
 	"github.com/ethereum/go-ethereum/cmd/utils"
 	"github.com/ethereum/go-ethereum/log"
 	cli "gopkg.in/urfave/cli.v1"
 )
 
-var (
-	gitCommit string // Git SHA1 commit hash of the release (set via linker flags)
+var gitCommit string // Git SHA1 commit hash of the release (set via linker flags)
+
+const (
+	defaultNodes   = 10
+	bzzServiceName = "bzz"
 )
 
-var (
-	topology  string
-	services  string
-	pivot     int
-	nodes     int
-	verbosity int
-)
+func main() {
+	err := newApp().Run(os.Args)
+	if err != nil {
+		log.Error(err.Error())
+		os.Exit(1)
+	}
+}
 
-var app = utils.NewApp("", "Swarm Snapshot Util")
-var discovery = true
-
-func init() {
-	log.PrintOrigins(true)
-	log.Root().SetHandler(log.LvlFilterHandler(log.Lvl(verbosity), log.StreamHandler(os.Stdout, log.TerminalFormat(true))))
+func newApp() (app *cli.App) {
+	app = utils.NewApp(gitCommit, "Swarm Snapshot Utility")
 
 	app.Name = "swarm-snapshot"
 	app.Usage = ""
 
 	app.Flags = []cli.Flag{
-		cli.StringFlag{
-			Name:        "topology",
-			Value:       "chain",
-			Usage:       "the desired topology to connect the nodes in (star, ring, chain, full)",
-			Destination: &topology,
-		},
-		cli.IntFlag{
-			Name:        "pivot",
-			Value:       0,
-			Usage:       "pivot node zero-index",
-			Destination: &pivot,
-		},
-		cli.IntFlag{
-			Name:        "nodes",
-			Value:       10,
-			Usage:       "swarm nodes",
-			Destination: &nodes,
-		},
 		cli.IntFlag{
 			Name:  "verbosity",
 			Value: 1,
-			Usage:       "verbosity",
-			Destination: &verbosity,
-		},
-		cli.StringFlag{
-			Name:        "services",
-			Value:       "",
-			Usage:       "comma separated list of services to boot the nodes with",
-			Destination: &services,
+			Usage: "verbosity level",
 		},
 	}
 
@@ -86,25 +59,27 @@ func init() {
 			Aliases: []string{"c"},
 			Usage:   "create a swarm snapshot",
 			Action:  create,
+			Flags: append(app.Flags,
+				cli.IntFlag{
+					Name:  "nodes",
+					Value: defaultNodes,
+					Usage: "number of nodes",
+				},
+				cli.StringFlag{
+					Name:  "services",
+					Value: bzzServiceName,
+					Usage: "comma separated list of services to boot the nodes with",
+				},
+			),
 		},
 		{
 			Name:    "verify",
 			Aliases: []string{"v"},
 			Usage:   "verify a swarm snapshot",
 			Action:  verify,
+			Flags:   app.Flags,
 		},
 	}
 
-	sort.Sort(cli.FlagsByName(app.Flags))
-	sort.Sort(cli.CommandsByName(app.Commands))
-}
-
-func main() {
-	err := app.Run(os.Args)
-	if err != nil {
-		log.Error(err.Error())
-
-		os.Exit(1)
-	}
-	os.Exit(0)
+	return app
 }

cmd/swarm/swarm-snapshot/run_test.go

@@ -27,7 +27,7 @@ import (
 
 func init() {
 	reexec.Register("swarm-snapshot", func() {
-		if err := app.Run(os.Args); err != nil {
+		if err := newApp().Run(os.Args); err != nil {
 			fmt.Fprintln(os.Stderr, err)
 			os.Exit(1)
 		}

cmd/swarm/swarm-snapshot/verify.go

@@ -19,10 +19,12 @@ package main
 import (
 	"context"
 	"errors"
+	"fmt"
+	"os"
 	"sync"
 	"time"
 
-	"github.com/ethereum/go-ethereum/cmd/utils"
+	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/node"
 	"github.com/ethereum/go-ethereum/p2p/simulations/adapters"
 	"github.com/ethereum/go-ethereum/swarm/network"
@@ -31,33 +33,23 @@ import (
 )
 
 func verify(ctx *cli.Context) error {
+	log.PrintOrigins(true)
+	log.Root().SetHandler(log.LvlFilterHandler(log.Lvl(ctx.Int("verbosity")), log.StreamHandler(os.Stdout, log.TerminalFormat(true))))
+
 	if len(ctx.Args()) < 1 {
-		return errors.New("argument should be the filename to verify or write-to")
+		return errors.New("argument should be the filename to verify")
 	}
-	filename, err := touchPath(ctx.Args()[0])
-	if err != nil {
-		return err
-	}
-	err = verifySnapshot(filename)
-	if err != nil {
-		utils.Fatalf("Simulation failed: %s", err)
+	return verifySnapshot(ctx.Args()[0])
 }
 
-	return err
-}
-
-func verifySnapshot(filename string) error {
+func verifySnapshot(filename string) (err error) {
 	sim := simulation.New(map[string]simulation.ServiceFunc{
-		"bzz": func(ctx *adapters.ServiceContext, b *sync.Map) (node.Service, func(), error) {
+		bzzServiceName: func(ctx *adapters.ServiceContext, b *sync.Map) (node.Service, func(), error) {
 			addr := network.NewAddr(ctx.Config.Node())
-
-			kp := network.NewKadParams()
-			kp.MinProxBinSize = testMinProxBinSize
-
-			kad := network.NewKademlia(addr.Over(), kp)
+			kad := network.NewKademlia(addr.Over(), network.NewKadParams())
 			hp := network.NewHiveParams()
 			hp.KeepAliveInterval = time.Duration(200) * time.Millisecond
-			hp.Discovery = true //discovery
+			hp.Discovery = false // discovery must be false when uploading a snapshot
 
 			config := &network.BzzConfig{
 				OverlayAddr:  addr.Over(),
@@ -65,20 +57,21 @@ func verifySnapshot(filename string) error {
 				HiveParams:   hp,
 			}
 			return network.NewBzz(config, kad, nil, nil, nil), nil, nil
-
 		},
 	})
 	defer sim.Close()
-	err := sim.UploadSnapshot(filename)
+
+	err = sim.UploadSnapshot(filename)
 	if err != nil {
-		utils.Fatalf("%v", err)
+		return fmt.Errorf("upload snapshot: %v", err)
 	}
 
 	ctx, cancelSimRun := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancelSimRun()
 
-	if _, err := sim.WaitTillHealthy(ctx, 2); err != nil {
-		utils.Fatalf("%v", err)
+	_, err = sim.WaitTillHealthy(ctx, 2)
+	if err != nil {
+		return fmt.Errorf("wait for healthy kademlia: %v", err)
 	}
 
 	return nil

cmd/swarm/swarm-snapshot/verify_test.go

@@ -22,7 +22,7 @@ import (
 
 func TestSnapshotVerify(t *testing.T) {
 	snap := runSnapshot(t,
-		"v",
+		"verify",
 		"testdata/snapshot.json",
 	)