Add unit tests

livepeer · Oct 31, 2024 · 41cfe06 · 41cfe06
1 parent d331225
commit 41cfe06
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 4 deletions.
diff --git a/server/realtime_payment.go b/server/realtime_payment.go
@@ -134,15 +134,15 @@ func (r *realtimePaymentSender) SendPayment(ctx context.Context, segmentInfo *Se
 }
 
 func (r *realtimePaymentReceiver) AccountPayment(
-	ctx context.Context, segmentInfo SegmentInfoReceiver) error {
+	ctx context.Context, segmentInfo *SegmentInfoReceiver) error {
 	fee := calculateFee(segmentInfo.inPixels, segmentInfo.priceInfo)
 
 	balance := r.orchestrator.Balance(segmentInfo.sender, core.ManifestID(segmentInfo.sessionID))
-	if balance.Cmp(fee) < 0 {
+	if balance == nil || balance.Cmp(fee) < 0 {
 		return errors.New("insufficient balance")
 	}
 	r.orchestrator.DebitFees(segmentInfo.sender, core.ManifestID(segmentInfo.sessionID), segmentInfo.priceInfo, segmentInfo.inPixels)
-	clog.V(common.DEBUG).Infof(ctx, "Accounted for payment for sessionID=%s, fee=%s", segmentInfo.sessionID, fee.FloatString(0))
+	clog.V(common.DEBUG).Infof(ctx, "Accounted payment for sessionID=%s, fee=%s", segmentInfo.sessionID, fee.FloatString(0))
 	return nil
 }
 

diff --git a/server/realtime_payment_test.go b/server/realtime_payment_test.go
@@ -80,3 +80,67 @@ func mockSender() pm.Sender {
 	sender.On("ValidateTicketParams", mock.Anything).Return(nil)
 	return sender
 }
+
+func TestAccountPayment(t *testing.T) {
+	require := require.New(t)
+
+	tests := []struct {
+		name   string
+		credit *big.Rat
+		expErr bool
+	}{
+		{
+			name:   "No credit",
+			credit: nil,
+			expErr: true,
+		},
+		{
+			name:   "Insufficient balance",
+			credit: new(big.Rat).SetInt64(900000),
+			expErr: true,
+		},
+		{
+			name:   "Sufficient balance",
+			credit: new(big.Rat).SetInt64(1100000),
+			expErr: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// given
+			sessionID := "abcdef"
+			sender := ethcommon.HexToAddress("0x0000000000000000000000000000000000000001")
+			segmentInfo := &SegmentInfoReceiver{
+				sender:    sender,
+				sessionID: sessionID,
+				inPixels:  1000000,
+				priceInfo: &net.PriceInfo{
+					PricePerUnit:  1,
+					PixelsPerUnit: 1,
+				},
+			}
+
+			node, _ := core.NewLivepeerNode(nil, "", nil)
+			balances := core.NewAddressBalances(1 * time.Minute)
+			node.Balances = balances
+			orch := core.NewOrchestrator(node, nil)
+
+			paymentReceiver := realtimePaymentReceiver{orchestrator: orch}
+			if tt.credit != nil {
+				node.Balances.Credit(sender, core.ManifestID(sessionID), tt.credit)
+			}
+
+			// when
+			err := paymentReceiver.AccountPayment(context.TODO(), segmentInfo)
+
+			// then
+			if tt.expErr {
+				require.Error(err, "insufficient balance")
+			} else {
+				require.Nil(err)
+			}
+		})
+	}
+
+}
diff --git a/server/segment_rpc.go b/server/segment_rpc.go
@@ -247,7 +247,7 @@ func startAccountLoop(orchestrator Orchestrator, sender ethcommon.Address, sessi
 		go func() {
 			for {
 				time.Sleep(5 * time.Second)
-				err := paymentReceiver.AccountPayment(context.TODO(), SegmentInfoReceiver{
+				err := paymentReceiver.AccountPayment(context.TODO(), &SegmentInfoReceiver{
 					sessionID: sessionID,
 					inPixels:  4000,
 					sender:    sender,