From e7b2f655965d0c8a320506af7539090e094faf66 Mon Sep 17 00:00:00 2001
From: Smaug123 <patrick+github@patrickstevens.co.uk>
Date: Sun, 6 Nov 2022 17:02:31 +0000
Subject: [PATCH] Add client data requests to generator

---
 Raft.Test/TestInMemoryServer.fs | 20 ++++++++++++--------
 Raft.Test/ValidHistory.fs       | 28 ++++++++++++++++++++++------
 2 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/Raft.Test/TestInMemoryServer.fs b/Raft.Test/TestInMemoryServer.fs
index c0e32f4..bca7b98 100644
--- a/Raft.Test/TestInMemoryServer.fs
+++ b/Raft.Test/TestInMemoryServer.fs
@@ -359,7 +359,9 @@ module TestInMemoryServer =
 
             firstTime = secondTime
 
-        property |> Prop.forAll (ValidHistory.arb clusterSize) |> check
+        property
+        |> Prop.forAll (ValidHistory.arb (Arb.Default.Byte().Generator) clusterSize)
+        |> check
 
 
     [<Test>]
@@ -380,7 +382,9 @@ module TestInMemoryServer =
 
             List.distinct leaders = leaders
 
-        property |> Prop.forAll (ValidHistory.arb clusterSize) |> check
+        property
+        |> Prop.forAll (ValidHistory.arb (Arb.Default.Byte().Generator) clusterSize)
+        |> check
 
     let duplicationProperty<'a when 'a : equality>
         (clusterSize : int)
@@ -432,19 +436,19 @@ module TestInMemoryServer =
             )
         )
 
-    let rec withDuplicateGen<'a> (clusterSize : int) : Gen<ValidHistory<'a> * ValidHistory<'a>> =
+    let rec withDuplicateGen<'a> (elementGen : Gen<'a>) (clusterSize : int) : Gen<ValidHistory<'a> * ValidHistory<'a>> =
         gen {
-            let! history = ValidHistory.gen clusterSize
+            let! history = ValidHistory.gen elementGen clusterSize
             let allDuplicatedHistories = allDuplicatedHistories<'a> clusterSize history
 
             match allDuplicatedHistories with
-            | [] -> return! withDuplicateGen clusterSize
+            | [] -> return! withDuplicateGen elementGen clusterSize
             | x -> return! Gen.elements x
         }
 
-    let duplicationArb<'a> (clusterSize : int) : Arbitrary<ValidHistory<'a> * ValidHistory<'a>> =
+    let duplicationArb<'a> (elementGen : Gen<'a>) (clusterSize : int) : Arbitrary<ValidHistory<'a> * ValidHistory<'a>> =
         { new Arbitrary<_>() with
-            member _.Generator = withDuplicateGen<'a> clusterSize
+            member _.Generator = withDuplicateGen<'a> elementGen clusterSize
 
             member _.Shrinker ((before, _withDuplicate)) =
                 ValidHistory.shrink<'a> clusterSize before
@@ -458,7 +462,7 @@ module TestInMemoryServer =
         let clusterSize = 5
 
         duplicationProperty<byte> clusterSize
-        |> Prop.forAll (duplicationArb clusterSize)
+        |> Prop.forAll (duplicationArb (Arb.Default.Byte().Generator) clusterSize)
         |> check
 
     [<Test>]
diff --git a/Raft.Test/ValidHistory.fs b/Raft.Test/ValidHistory.fs
index 518b8d1..ff27eac 100644
--- a/Raft.Test/ValidHistory.fs
+++ b/Raft.Test/ValidHistory.fs
@@ -34,7 +34,12 @@ module ValidHistory =
             | NetworkAction.ClientRequest _ -> return! networkActionGenNoClientRequests clusterSize
         }
 
-    let private historyGenOfLength<'a> (clusterSize : int) (len : int) : Gen<NetworkAction<'a> list> =
+    let private historyGenOfLength<'a>
+        (elementGen : Gen<'a>)
+        (clusterSize : int)
+        (len : int)
+        : Gen<NetworkAction<'a> list>
+        =
         let cluster, network = InMemoryCluster.make<'a> clusterSize
         // Note: takes a reversed list.
         let permissibleNext () : NetworkAction<'a> list =
@@ -52,21 +57,32 @@ module ValidHistory =
                     yield NetworkAction.InactivityTimeout server
             ]
 
+        let clientRequestGen =
+            gen {
+                let! element = elementGen
+                let! id = Gen.choose (0, clusterSize - 1)
+                return NetworkAction.ClientRequest (id * 1<ServerId>, element)
+            }
+
         let rec go (len : int) =
             gen {
                 if len = 0 then
                     return []
                 else
                     let! smaller = go (len - 1)
-                    let! next = Gen.elements (permissibleNext ())
+
+                    let! next =
+                        clientRequestGen :: List.replicate 5 (Gen.elements (permissibleNext ()))
+                        |> Gen.oneof
+
                     NetworkAction.perform cluster network next
                     return next :: smaller
             }
 
         go (abs len)
 
-    let gen<'a> (clusterSize : int) : Gen<ValidHistory<'a>> =
-        historyGenOfLength<'a> clusterSize
+    let gen<'a> (elementGen : Gen<'a>) (clusterSize : int) : Gen<ValidHistory<'a>> =
+        historyGenOfLength<'a> elementGen clusterSize
         |> Gen.sized
         |> Gen.map (List.rev >> ValidHistory)
 
@@ -100,9 +116,9 @@ module ValidHistory =
 
         Seq.concat [ removeOne ; shrinkMessageId ]
 
-    let arb<'a> (clusterSize : int) =
+    let arb<'a> (elementGen : Gen<'a>) (clusterSize : int) =
         { new Arbitrary<ValidHistory<'a>>() with
-            override _.Generator = gen clusterSize
+            override _.Generator = gen elementGen clusterSize
 
             override _.Shrinker history = shrink clusterSize history
         }