simplify api even more, make it also more robust

1 files changed, 26 insertions, 14 deletions

diff --git a/nizk/stage2_test.go b/nizk/stage2_test.go
index e5b44c6..446f20d 100644
--- a/nizk/stage2_test.go
+++ b/nizk/stage2_test.go
@@ -7,21 +7,33 @@ import (
 )
 
 func TestStage2Simple(t *testing.T) {
-	id1 := Curve.RandomScalar()
-	b1 := NewBit(id1, false)
-	b2 := NewBit(id1, true)
-
+	id := Curve.RandomScalar()
+	b1, _, _ := NewBit(id, false) // This is also the junction
+	r1, _ := b1.RevealStage1()
 	c1 := b1.StageCommit()
-	r1, _, _ := b1.RevealStage1()
 
-	bc2, _ := b2.Commit()
-	c2 := b2.StageCommit()
-	r2, p2, e := b2.RevealStage2(b1)
-	if e != nil {
-		t.Fatalf("e: %v", e)
-	}
-	if !bc2.VerifyStage2(c1, c2, r1, r2, p2) {
-		t.Fatalf("failed to verify!\nbc2: %#v\nc1: %#v\nc2: %#v\nr1: %#v\nr2: %#v\np2: %#v\n",
-			bc2, c1, c2, r1, r2, p2)
+	for _, s := range [][2]bool{
+		{false, false},
+		{true, false},
+		{true, true},
+		{false, true},
+	} {
+		b2, bc2, _ := NewBit(id, s[0])
+		b3, bc3, _ := NewBit(id, s[1])
+
+		c2 := b2.StageCommit()
+		c3 := b3.StageCommit()
+
+		r2, p2 := b2.RevealStage2(true, b1)
+		if !bc2.VerifyStage2(c1, c2, r1, r2, p2) {
+			t.Fatalf("failed to verify!\nbc2: %#v\nc1: %#v\nc2: %#v\nr1: %#v\nr2: %#v\np2: %#v\n",
+				bc2, c1, c2, r1, r2, p2)
+		}
+
+		r3, p3 := b3.RevealStage2(true, b1)
+		if !bc3.VerifyStage2(c1, c3, r1, r3, p3) {
+			t.Fatalf("faild to verify bc3")
+		}
 	}
+
 }