From 2bbc7d949f844558cb2fb08abf0e0d4a078efc5b Mon Sep 17 00:00:00 2001
From: Jernej Kos <jernej@kos.mx>
Date: Tue, 21 Feb 2023 20:08:10 +0100
Subject: [PATCH 1/2] tests: Cleanup txsource temporary identity directories

---
 .../cmd/debug/txsource/workload/registration.go     | 13 ++++++++-----
 .../cmd/debug/txsource/workload/workload.go         |  6 ------
 .../scenario/e2e/runtime/txsource.go                |  1 -
 3 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/go/oasis-node/cmd/debug/txsource/workload/registration.go b/go/oasis-node/cmd/debug/txsource/workload/registration.go
index 7991c9cfc5b..66a0e5a33ca 100644
--- a/go/oasis-node/cmd/debug/txsource/workload/registration.go
+++ b/go/oasis-node/cmd/debug/txsource/workload/registration.go
@@ -6,10 +6,8 @@ import (
 	"math/rand"
 	"net"
 	"os"
-	"path/filepath"
 	"time"
 
-	"github.com/spf13/viper"
 	"google.golang.org/grpc"
 
 	beacon "github.com/oasisprotocol/oasis-core/go/beacon/api"
@@ -181,11 +179,11 @@ func (r *registration) Run( // nolint: gocyclo
 		panic(err)
 	}
 
-	baseDir := viper.GetString(CfgDataDir)
-	nodeIdentitiesDir := filepath.Join(baseDir, "node-identities")
-	if err = common.Mkdir(nodeIdentitiesDir); err != nil {
+	nodeIdentitiesDir, err := os.MkdirTemp("", "oasis-e2e-registration")
+	if err != nil {
 		return fmt.Errorf("txsource/registration: failed to create node-identities dir: %w", err)
 	}
+	defer os.RemoveAll(nodeIdentitiesDir)
 
 	type runtimeInfo struct {
 		entityIdx int
@@ -256,6 +254,9 @@ func (r *registration) Run( // nolint: gocyclo
 
 			entityAccs[i].nodeIdentities = append(entityAccs[i].nodeIdentities, &nodeAcc{ident, nodeDesc, nodeAccNonce})
 			ent.Nodes = append(ent.Nodes, ident.NodeSigner.Public())
+
+			// Cleanup temporary node identity directory after generation.
+			_ = os.RemoveAll(dataDir)
 		}
 
 		// Register entity.
@@ -302,6 +303,8 @@ func (r *registration) Run( // nolint: gocyclo
 			}
 		}
 	}
+	// Cleanup temporary identities directory after generation.
+	_ = os.RemoveAll(nodeIdentitiesDir)
 
 	iteration := 0
 	var loopCtx context.Context
diff --git a/go/oasis-node/cmd/debug/txsource/workload/workload.go b/go/oasis-node/cmd/debug/txsource/workload/workload.go
index 02c87a15325..4740533c358 100644
--- a/go/oasis-node/cmd/debug/txsource/workload/workload.go
+++ b/go/oasis-node/cmd/debug/txsource/workload/workload.go
@@ -7,7 +7,6 @@ import (
 	"time"
 
 	flag "github.com/spf13/pflag"
-	"github.com/spf13/viper"
 	"google.golang.org/grpc"
 
 	"github.com/oasisprotocol/oasis-core/go/common/crypto/signature"
@@ -20,8 +19,6 @@ import (
 )
 
 const (
-	CfgDataDir = "datadir"
-
 	maxSubmissionRetryElapsedTime = 120 * time.Second
 
 	fundAccountAmount = 10000000000
@@ -197,9 +194,6 @@ func FundAccountFromTestEntity(
 }
 
 func init() {
-	Flags.String(CfgDataDir, "", "data directory")
-	_ = viper.BindPFlags(Flags)
-
 	Flags.AddFlagSet(QueriesFlags)
 	Flags.AddFlagSet(RuntimeFlags)
 }
diff --git a/go/oasis-test-runner/scenario/e2e/runtime/txsource.go b/go/oasis-test-runner/scenario/e2e/runtime/txsource.go
index e60079a999b..9e89a12379a 100644
--- a/go/oasis-test-runner/scenario/e2e/runtime/txsource.go
+++ b/go/oasis-test-runner/scenario/e2e/runtime/txsource.go
@@ -789,7 +789,6 @@ func (sc *txSourceImpl) startWorkload(childEnv *env.Env, errCh chan error, name
 		"debug", "txsource",
 		"--address", "unix:" + node.SocketPath(),
 		"--" + common.CfgDebugAllowTestKeys,
-		"--" + workload.CfgDataDir, d.String(),
 		"--" + flags.CfgDebugDontBlameOasis,
 		"--" + flags.CfgDebugTestEntity,
 		"--" + commonGrpc.CfgLogDebug,

From cb9ded1b181e5d65c6e2721f2a89a976b76c59b7 Mon Sep 17 00:00:00 2001
From: Jernej Kos <jernej@kos.mx>
Date: Tue, 21 Feb 2023 20:11:56 +0100
Subject: [PATCH 2/2] tests: Ensure node is synced before doing further queries

---
 .changelog/5192.trivial.md                           | 0
 go/oasis-node/cmd/debug/txsource/workload/queries.go | 8 ++++++++
 2 files changed, 8 insertions(+)
 create mode 100644 .changelog/5192.trivial.md

diff --git a/.changelog/5192.trivial.md b/.changelog/5192.trivial.md
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/go/oasis-node/cmd/debug/txsource/workload/queries.go b/go/oasis-node/cmd/debug/txsource/workload/queries.go
index d6f4a2b8bc6..b5755dde5c5 100644
--- a/go/oasis-node/cmd/debug/txsource/workload/queries.go
+++ b/go/oasis-node/cmd/debug/txsource/workload/queries.go
@@ -872,6 +872,14 @@ func (q *queries) Run(
 	for {
 		loopCtx, cancel := context.WithTimeout(ctx, queriesIterationTimeout)
 
+		// Ensure the node appears synced before doing queries. Given that nodes can be killed and
+		// held stopped for a while, they need to sync once they come back.
+		if isSynced, _ := q.control.IsSynced(loopCtx); !isSynced {
+			_ = q.control.WaitSync(loopCtx)
+			time.Sleep(1 * time.Second)
+			continue
+		}
+
 		err := q.doQueries(loopCtx, rng)
 		cancel()
 		switch {