swarm/simulation: UploadSnapshot updated

swarm/network/simulation/kademlia.go

@@ -20,9 +20,6 @@ import (
 	"context"
 	"encoding/binary"
 	"encoding/hex"
-	"encoding/json"
-	"io/ioutil"
-	"os"
 	"time"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -204,21 +201,3 @@ func removeDuplicatesAndSingletons(arr []uint64) []uint64 {
 
 	return arr
 }
-
-func ReadSnapshot(filename string) (*simulations.Snapshot, error) {
-	f, err := os.Open(filename)
-	if err != nil {
-		return nil, err
-	}
-	defer f.Close()
-	jsonbyte, err := ioutil.ReadAll(f)
-	if err != nil {
-		return nil, err
-	}
-	var snap simulations.Snapshot
-	err = json.Unmarshal(jsonbyte, &snap)
-	if err != nil {
-		return nil, err
-	}
-	return &snap, nil
-}

swarm/network/simulation/node.go

@@ -17,6 +17,7 @@
 package simulation
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"io/ioutil"
@@ -217,7 +218,7 @@ func (s *Simulation) AddNodesAndConnectStar(count int, opts ...AddNodeOption) (i
 // UploadSnapshot uploads a snapshot to the simulation
 // This method tries to open the json file provided, applies the config to all nodes
 // and then loads the snapshot into the Simulation network
-func (s *Simulation) UploadSnapshot(snapshotFile string, opts ...AddNodeOption) error {
+func (s *Simulation) UploadSnapshot(ctx context.Context, snapshotFile string, opts ...AddNodeOption) error {
 	f, err := os.Open(snapshotFile)
 	if err != nil {
 		return err
@@ -256,8 +257,14 @@ func (s *Simulation) UploadSnapshot(snapshotFile string, opts ...AddNodeOption)
 	if err != nil {
 		return err
 	}
+
+	err = s.WaitTillSnapshotRecreated(ctx, &snap)
+	if err == nil {
 		log.Info("Snapshot loaded")
-	return nil
+	} else {
+		log.Warn("Snapshot load failed", "error", err.Error())
+	}
+	return err
 }
 
 // StartNode starts a node by NodeID.

swarm/network/simulation/node_test.go

@@ -289,6 +289,7 @@ func TestUploadSnapshot(t *testing.T) {
 				HiveParams:   hp,
 			}
 			kad := network.NewKademlia(addr.Over(), network.NewKadParams())
+			b.Store(BucketKeyKademlia, kad)
 			return network.NewBzz(config, kad, nil, nil, nil), nil, nil
 		},
 	})
@@ -296,12 +297,13 @@ func TestUploadSnapshot(t *testing.T) {
 
 	nodeCount := 16
 	log.Debug("Uploading snapshot")
-	err := s.UploadSnapshot(fmt.Sprintf("../stream/testing/snapshot_%d.json", nodeCount))
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
+	defer cancel()
+	err := s.UploadSnapshot(ctx, fmt.Sprintf("../stream/testing/snapshot_%d.json", nodeCount))
 	if err != nil {
 		t.Fatalf("Error uploading snapshot to simulation network: %v", err)
 	}
 
-	ctx := context.Background()
 	log.Debug("Starting simulation...")
 	s.Run(ctx, func(ctx context.Context, sim *Simulation) error {
 		log.Debug("Checking")

swarm/network/stream/snapshot_retrieval_test.go

@@ -154,15 +154,15 @@ func runFileRetrievalTest(nodeCount int) error {
 	//array where the generated chunk hashes will be stored
 	conf.hashes = make([]storage.Address, 0)
 
+	ctx, cancelSimRun := context.WithTimeout(context.Background(), 3*time.Minute)
+	defer cancelSimRun()
+
 	filename := fmt.Sprintf("testing/snapshot_%d.json", nodeCount)
-	err := sim.UploadSnapshot(filename)
+	err := sim.UploadSnapshot(ctx, filename)
 	if err != nil {
 		return err
 	}
 
-	ctx, cancelSimRun := context.WithTimeout(context.Background(), 3*time.Minute)
-	defer cancelSimRun()
-
 	log.Info("Starting simulation")
 
 	result := sim.Run(ctx, func(ctx context.Context, sim *simulation.Simulation) error {
@@ -188,13 +188,6 @@ func runFileRetrievalTest(nodeCount int) error {
 		if err != nil {
 			return err
 		}
-		snap, err := simulation.ReadSnapshot(filename)
-		if err != nil {
-			return err
-		}
-		if err := sim.WaitTillSnapshotRecreated(ctx, snap); err != nil {
-			return err
-		}
 
 		log.Info("network healthy, start file checks")
 
@@ -257,13 +250,15 @@ func runRetrievalTest(t *testing.T, chunkCount int, nodeCount int) error {
 	//array where the generated chunk hashes will be stored
 	conf.hashes = make([]storage.Address, 0)
 
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
+	defer cancel()
+
 	filename := fmt.Sprintf("testing/snapshot_%d.json", nodeCount)
-	err := sim.UploadSnapshot(filename)
+	err := sim.UploadSnapshot(ctx, filename)
 	if err != nil {
 		return err
 	}
 
-	ctx := context.Background()
 	result := sim.Run(ctx, func(ctx context.Context, sim *simulation.Simulation) error {
 		nodeIDs := sim.UpNodeIDs()
 		for _, n := range nodeIDs {
@@ -288,13 +283,6 @@ func runRetrievalTest(t *testing.T, chunkCount int, nodeCount int) error {
 		if err != nil {
 			return err
 		}
-		snap, err := simulation.ReadSnapshot(filename)
-		if err != nil {
-			t.Fatalf("failed to read snapshot: %s", err)
-		}
-		if err := sim.WaitTillSnapshotRecreated(ctx, snap); err != nil {
-			return err
-		}
 
 		// File retrieval check is repeated until all uploaded files are retrieved from all nodes
 		// or until the timeout is reached.

swarm/network/stream/snapshot_sync_test.go

@@ -147,25 +147,16 @@ func testSyncingViaGlobalSync(t *testing.T, chunkCount int, nodeCount int) {
 	//array where the generated chunk hashes will be stored
 	conf.hashes = make([]storage.Address, 0)
 
-	filename := fmt.Sprintf("testing/snapshot_%d.json", nodeCount)
+	ctx, cancelSimRun := context.WithTimeout(context.Background(), 3*time.Minute)
-	err := sim.UploadSnapshot(filename)
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	ctx, cancelSimRun := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer cancelSimRun()
 
-	snap, err := simulation.ReadSnapshot(filename)
+	filename := fmt.Sprintf("testing/snapshot_%d.json", nodeCount)
+	err := sim.UploadSnapshot(ctx, filename)
 	if err != nil {
-		t.Fatalf("failed to read snapshot: %s", err)
-	}
-	if err := sim.WaitTillSnapshotRecreated(ctx, snap); err != nil {
 		t.Fatal(err)
 	}
 
 	result := runSim(conf, ctx, sim, chunkCount)
-
 	if result.Error != nil {
 		t.Fatal(result.Error)
 	}

swarm/network/stream/streamer_test.go

@@ -1257,8 +1257,9 @@ func TestGetSubscriptionsRPC(t *testing.T) {
 		simulation.NewPeerEventsFilter().ReceivedMessages().Protocol("stream").MsgCode(subscribeMsgCode),
 	)
 
-	// upload a snapshot
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*120)
-	err := sim.UploadSnapshot(fmt.Sprintf("testing/snapshot_%d.json", nodeCount))
+	defer cancel()
+	err := sim.UploadSnapshot(ctx, fmt.Sprintf("testing/snapshot_%d.json", nodeCount))
 	if err != nil {
 		t.Fatal(err)
 	}

swarm/pss/prox_test.go

@@ -213,20 +213,12 @@ func testProxNetwork(t *testing.T) {
 	services := newProxServices(tstdata, true, handlerContextFuncs, tstdata.kademlias)
 	tstdata.sim = simulation.New(services)
 	defer tstdata.sim.Close()
-	filename := fmt.Sprintf("testdata/snapshot_%d.json", nodeCount)
-	err := tstdata.sim.UploadSnapshot(filename)
-	if err != nil {
-		t.Fatal(err)
-	}
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*120)
 	defer cancel()
-	snap, err := simulation.ReadSnapshot(filename)
+	filename := fmt.Sprintf("testdata/snapshot_%d.json", nodeCount)
+	err := tstdata.sim.UploadSnapshot(ctx, filename)
 	if err != nil {
-		t.Fatalf("failed to read snapshot: %s", err)
+		t.Fatal(err)
-	}
-	err = tstdata.sim.WaitTillSnapshotRecreated(ctx, snap)
-	if err != nil {
-		t.Fatalf("failed to recreate snapshot: %s", err)
 	}
 	tstdata.init(msgCount) // initialize the test data
 	wrapper := func(c context.Context, _ *simulation.Simulation) error {