From 29479d21a014b9d1d96de98a08a27a2713b7399c Mon Sep 17 00:00:00 2001
From: Nikolay Eskov <mr.eskov1@yandex.ru>
Date: Thu, 17 Oct 2024 03:41:17 +0300
Subject: [PATCH] Add
 'TestUpdateGeneratingBalanceDuringScriptExecutionAfterLightNode' test.

---
 pkg/ride/environment.go          |  6 +--
 pkg/ride/test_helpers_test.go    | 14 ++----
 pkg/ride/tree_evaluation_test.go | 78 ++++++++++++++++++++++++++++++++
 3 files changed, 86 insertions(+), 12 deletions(-)

diff --git a/pkg/ride/environment.go b/pkg/ride/environment.go
index 49c4ff8aa..0d9eaf991 100644
--- a/pkg/ride/environment.go
+++ b/pkg/ride/environment.go
@@ -32,14 +32,14 @@ type WrappedState struct {
 }
 
 func newWrappedState(
-	env *EvaluationEnvironment,
+	env environment,
 	originalStateForWrappedState types.EnrichedSmartState,
 	rootScriptLibVersion ast.LibraryVersion,
 ) *WrappedState {
 	return &WrappedState{
 		diff:                      newDiffState(originalStateForWrappedState),
-		cle:                       env.th.(rideAddress),
-		scheme:                    env.sch,
+		cle:                       env.this().(rideAddress),
+		scheme:                    env.scheme(),
 		height:                    proto.Height(env.height()),
 		rootScriptLibVersion:      rootScriptLibVersion,
 		rootActionsCountValidator: proto.NewScriptActionsCountValidator(),
diff --git a/pkg/ride/test_helpers_test.go b/pkg/ride/test_helpers_test.go
index 32d003798..ba0242e64 100644
--- a/pkg/ride/test_helpers_test.go
+++ b/pkg/ride/test_helpers_test.go
@@ -635,16 +635,12 @@ func (e *testEnv) withDataFromJSON(s string) *testEnv {
 
 func (e *testEnv) withWrappedState() *testEnv {
 	v, err := e.me.libVersion()
-	if err != nil {
-		panic(err)
-	}
-	e.ws = &WrappedState{
-		diff:                      newDiffState(e.ms),
-		cle:                       e.me.this().(rideAddress),
-		scheme:                    e.me.scheme(),
-		rootScriptLibVersion:      v,
-		rootActionsCountValidator: proto.NewScriptActionsCountValidator(),
+	require.NoError(e.t, err)
+	if e.me.heightFunc == nil { // create stub height function`
+		e.me.heightFunc = func() rideInt { return 0 }
+		defer func() { e.me.heightFunc = nil }()
 	}
+	e.ws = newWrappedState(e.me, e.ms, v)
 	e.me.stateFunc = func() types.SmartState {
 		return e.ws
 	}
diff --git a/pkg/ride/tree_evaluation_test.go b/pkg/ride/tree_evaluation_test.go
index 8c6fa1b0a..66c03d038 100644
--- a/pkg/ride/tree_evaluation_test.go
+++ b/pkg/ride/tree_evaluation_test.go
@@ -6330,3 +6330,81 @@ func TestAttachedPaymentsValidation(t *testing.T) {
 		})
 	})
 }
+
+func TestUpdateGeneratingBalanceDuringScriptExecutionAfterLightNode(t *testing.T) {
+	sender := newTestAccount(t, "SENDER") // 3N8CkZAyS4XcDoJTJoKNuNk2xmNKmQj7myW
+	dApp1 := newTestAccount(t, "DAPP1")   // 3MzDtgL5yw73C2xVLnLJCrT5gCL4357a4sz
+
+	const src = `
+		{-# STDLIB_VERSION 5 #-}
+		{-# CONTENT_TYPE DAPP #-}
+		{-# SCRIPT_TYPE ACCOUNT #-}
+		
+		@Callable(i)
+		func call(to: String, value: Int) = {
+		  let balanceBeforeTransfer = wavesBalance(this)
+		  let balanceAfter = invoke(this, "transfer", [to, value], [])
+		  let balanceAfterTransfer = wavesBalance(this)
+		  let isUpdated = balanceAfter == balanceAfterTransfer.generating
+		  if (!isUpdated && balanceBeforeTransfer.generating != balanceAfterTransfer.generating) then
+			throw("failure")
+		  else
+		  ([], isUpdated)
+		}
+		
+		@Callable(i)
+		func transfer(to: String, value: Int) = {
+		  let balance = wavesBalance(this)
+		  let toAddr = addressFromStringValue(to)
+		  ([ScriptTransfer(toAddr, value, unit)], balance.generating - value)
+		}`
+	tree, errs := ridec.CompileToTree(src)
+	require.Empty(t, errs)
+
+	const (
+		dAppInitialBalance = 10 * proto.PriceConstant
+		transferAmount     = 5 * proto.PriceConstant
+	)
+
+	createEnv := func(t *testing.T, lightNodeActivated bool) *testEnv {
+		complexity, err := MaxChainInvokeComplexityByVersion(ast.LibV5)
+		require.NoError(t, err)
+		env := newTestEnv(t).withLibVersion(ast.LibV5).withComplexityLimit(int(complexity)).
+			withBlockV5Activated().withProtobufTx().
+			withDataEntriesSizeV2().withMessageLengthV3().withValidateInternalPayments().
+			withThis(dApp1).withDApp(dApp1).withSender(sender).
+			withInvocation("call", withTransactionID(crypto.Digest{})).withTree(dApp1, tree).
+			withWavesBalance(sender, 0).
+			withWavesBalance(dApp1, dAppInitialBalance, 0, 0, dAppInitialBalance)
+		if lightNodeActivated {
+			env = env.withLightNodeActivated()
+		}
+		return env.withWrappedState()
+	}
+	t.Run("before-light-node", func(t *testing.T) {
+		env := createEnv(t, false).toEnv()
+		fc := proto.NewFunctionCall("call", proto.Arguments{
+			proto.NewStringArgument(sender.address().String()),
+			proto.NewIntegerArgument(transferAmount),
+		})
+		res, err := CallFunction(env, tree, fc)
+		require.NoError(t, err)
+		require.Len(t, res.ScriptActions(), 1)
+		isBalanceUpdated, ok := res.userResult().(rideBoolean)
+		require.True(t, ok)
+		assert.False(t, bool(isBalanceUpdated))
+	})
+	t.Run("after-light-node", func(t *testing.T) {
+		env := createEnv(t, true).toEnv()
+		fc := proto.NewFunctionCall("call", proto.Arguments{
+			proto.NewStringArgument(sender.address().String()),
+			proto.NewIntegerArgument(transferAmount),
+		})
+		res, err := CallFunction(env, tree, fc)
+		require.NoError(t, err)
+		require.Len(t, res.ScriptActions(), 1)
+		isBalanceUpdated, ok := res.userResult().(rideBoolean)
+		require.True(t, ok)
+		assert.True(t, bool(isBalanceUpdated))
+	})
+}