From 2dc46af4155a74420e84ba0523a63b00dea3fb4e Mon Sep 17 00:00:00 2001
From: vyzo <vyzo@hackzen.org>
Date: Fri, 8 May 2020 20:07:05 +0300
Subject: [PATCH] update tests for quadratic p4

---
 gossipsub_spam_test.go |  2 +-
 score_test.go          | 15 ++++++++-------
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/gossipsub_spam_test.go b/gossipsub_spam_test.go
index 387a9e5..e92248e 100644
--- a/gossipsub_spam_test.go
+++ b/gossipsub_spam_test.go
@@ -669,7 +669,7 @@ func TestGossipsubAttackInvalidMessageSpam(t *testing.T) {
 					time.Sleep(100*time.Millisecond + GossipSubHeartbeatInitialDelay)
 
 					// The attackers score should now have fallen below zero
-					if attackerScore() > 0 {
+					if attackerScore() >= 0 {
 						t.Fatalf("Expected attacker score to be less than zero but it's %f", attackerScore())
 					}
 					// There should be several rejected messages (because the signature was invalid)
diff --git a/score_test.go b/score_test.go
index aef53ad..82d98fd 100644
--- a/score_test.go
+++ b/score_test.go
@@ -1,6 +1,7 @@
 package pubsub
 
 import (
+	"math"
 	"sync"
 	"testing"
 	"time"
@@ -457,7 +458,7 @@ func TestScoreInvalidMessageDeliveries(t *testing.T) {
 	topicScoreParams := &TopicScoreParams{
 		TopicWeight:                    1,
 		TimeInMeshQuantum:              time.Second,
-		InvalidMessageDeliveriesWeight: 1,
+		InvalidMessageDeliveriesWeight: -1,
 		InvalidMessageDeliveriesDecay:  1.0,
 	}
 	params.Topics[mytopic] = topicScoreParams
@@ -478,7 +479,7 @@ func TestScoreInvalidMessageDeliveries(t *testing.T) {
 
 	ps.refreshScores()
 	aScore := ps.Score(peerA)
-	expected := topicScoreParams.TopicWeight * topicScoreParams.InvalidMessageDeliveriesWeight * float64(nMessages)
+	expected := topicScoreParams.TopicWeight * topicScoreParams.InvalidMessageDeliveriesWeight * float64(nMessages*nMessages)
 	if aScore != expected {
 		t.Fatalf("Score: %f. Expected %f", aScore, expected)
 	}
@@ -494,7 +495,7 @@ func TestScoreInvalidMessageDeliveriesDecay(t *testing.T) {
 	topicScoreParams := &TopicScoreParams{
 		TopicWeight:                    1,
 		TimeInMeshQuantum:              time.Second,
-		InvalidMessageDeliveriesWeight: 1,
+		InvalidMessageDeliveriesWeight: -1,
 		InvalidMessageDeliveriesDecay:  0.9,
 	}
 	params.Topics[mytopic] = topicScoreParams
@@ -515,7 +516,7 @@ func TestScoreInvalidMessageDeliveriesDecay(t *testing.T) {
 
 	ps.refreshScores()
 	aScore := ps.Score(peerA)
-	expected := topicScoreParams.TopicWeight * topicScoreParams.InvalidMessageDeliveriesWeight * topicScoreParams.InvalidMessageDeliveriesDecay * float64(nMessages)
+	expected := topicScoreParams.TopicWeight * topicScoreParams.InvalidMessageDeliveriesWeight * math.Pow(topicScoreParams.InvalidMessageDeliveriesDecay*float64(nMessages), 2)
 	if aScore != expected {
 		t.Fatalf("Score: %f. Expected %f", aScore, expected)
 	}
@@ -523,7 +524,7 @@ func TestScoreInvalidMessageDeliveriesDecay(t *testing.T) {
 	// refresh scores a few times to apply decay
 	for i := 0; i < 10; i++ {
 		ps.refreshScores()
-		expected *= topicScoreParams.InvalidMessageDeliveriesDecay
+		expected *= math.Pow(topicScoreParams.InvalidMessageDeliveriesDecay, 2)
 	}
 	aScore = ps.Score(peerA)
 	if aScore != expected {
@@ -651,13 +652,13 @@ func TestScoreRejectMessageDeliveries(t *testing.T) {
 	ps.RejectMessage(&msg, rejectValidationFailed)
 
 	aScore = ps.Score(peerA)
-	expected = -2.0
+	expected = -4.0
 	if aScore != expected {
 		t.Fatalf("Score: %f. Expected %f", aScore, expected)
 	}
 
 	bScore = ps.Score(peerB)
-	expected = -2.0
+	expected = -4.0
 	if bScore != expected {
 		t.Fatalf("Score: %f. Expected %f", bScore, expected)
 	}