Merge pull request #419 from JohnLCaron/consumerTests

Bring proto and json consumer up to par (almost).

egklib/src/commonMain/kotlin/electionguard/publish/ElectionRecordJsonPaths.kt

@@ -53,9 +53,9 @@ data class ElectionRecordJsonPaths(val topDir : String) {
         return "$ballotDir/$PLAINTEXT_BALLOT_PREFIX$id$JSON_SUFFIX"
     }
 
-    fun encryptedBallotPath(outputDir : String, ballotId : String): String {
+    fun encryptedBallotPath(ballotDir : String, ballotId : String): String {
         val id = ballotId.replace(" ", "_")
-        return "${outputDir}/$ENCRYPTED_BALLOT_PREFIX$id$JSON_SUFFIX"
+        return "${ballotDir}/$ENCRYPTED_BALLOT_PREFIX$id$JSON_SUFFIX"
     }
 
     fun pepBallotPath(outputDir : String, ballotId : String): String {

egklib/src/commonMain/kotlin/electionguard/publish/ElectionRecordProtoPaths.kt

@@ -70,7 +70,12 @@ data class ElectionRecordProtoPaths(val topDir : String) {
         return "$electionRecordDir/$ENCRYPTED_DIR/$useDevice/"
     }
 
-    fun encryptedBallotPath(device: String, ballotId: String): String {
+    fun encryptedBallotPath(ballotDir : String, ballotId : String): String {
+        val id = ballotId.replace(" ", "_")
+        return "${ballotDir}/${ENCRYPTED_BALLOT_PREFIX}$id${PROTO_SUFFIX}"
+    }
+
+    fun encryptedBallotDevicePath(device: String, ballotId: String): String {
         val useDevice = device.replace(" ", "_")
         val id = ballotId.replace(" ", "_")
         return "$electionRecordDir/$ENCRYPTED_DIR/$useDevice/${ENCRYPTED_BALLOT_PREFIX}$id${PROTO_SUFFIX}"

egklib/src/commonTest/kotlin/electionguard/publish/ConsumerTest.kt

@@ -1,11 +1,14 @@
 package electionguard.publish
 
 import com.github.michaelbull.result.*
+import electionguard.ballot.EncryptedBallot
 import electionguard.core.productionGroup
 import electionguard.core.runTest
 import electionguard.decrypt.DecryptingTrusteeIF
+import electionguard.util.ErrorMessages
 import org.junit.jupiter.api.Assertions.assertNotNull
 import kotlin.test.Test
+import kotlin.test.assertEquals
 import kotlin.test.assertFalse
 import kotlin.test.assertTrue
 
@@ -16,13 +19,16 @@ class ConsumerTest {
     @Test
     fun consumerJson() {
         val consumerIn = makeConsumer(group, topdir)
+        assertTrue(consumerIn.isJson())
         assertTrue(consumerIn.readElectionConfig() is Ok) // proto doesnt always have config, may have just init
         testConsumer(consumerIn)
     }
 
     @Test
     fun consumerProto() {
-        testConsumer(makeConsumer(group, "src/commonTest/data/workflow/chainedProto"))
+        val consumerIn = makeConsumer(group, "src/commonTest/data/workflow/chainedProto")
+        assertFalse(consumerIn.isJson())
+        testConsumer(consumerIn)
     }
 
     fun testConsumer(consumerIn : Consumer) {
@@ -35,17 +41,37 @@ class ConsumerTest {
 
         assertTrue( consumerIn.readEncryptedBallotChain(device) is Ok)
 
-        var iter = consumerIn.iterateEncryptedBallots(device) { true }
-        assertNotNull(iter)
-        assertTrue(iter.count() > 0)
+        val iterb = consumerIn.iterateEncryptedBallots(device) { true }
+        assertNotNull(iterb)
+        val iter = iterb.iterator()
+        var count = 0
+        var eballot : EncryptedBallot? = null
+        while (iter.hasNext()) {
+            eballot = iter.next()
+            assertTrue(eballot is EncryptedBallot)
+            count++
+        }
+        assertNotNull(eballot)
+        assertEquals(count, iterb.count())
 
-        iter = consumerIn.iterateAllEncryptedBallots { true }
+        val ballotDir = "${consumerIn.topdir()}/encrypted_ballots/$device"
+        val readResult = consumerIn.readEncryptedBallot(ballotDir, eballot!!.ballotId)
+        println(readResult)
+        assertTrue (readResult is Ok)
+        assertEquals(eballot, readResult.unwrap())
+
+        val iterAll = consumerIn.iterateAllEncryptedBallots { true }
         assertNotNull(iter)
-        assertTrue(iter.count() > 0)
+        assertTrue(iterAll.count() > 0)
 
         val iter2 = consumerIn.iterateDecryptedBallots()
         assertNotNull(iter2)
-        assertTrue(iter.count() >= 0)
+        assertTrue(iter2.count() > 0)
+
+        val plaintextDir = "${consumerIn.topdir()}/private_data/input"
+        val iter3 = consumerIn.iteratePlaintextBallots(plaintextDir) { true }
+        assertNotNull(iter3)
+        assertTrue(iter3.count() > 0)
     }
 
     @Test

egklib/src/jvmMain/kotlin/electionguard/publish/ConsumerJson.kt

@@ -156,11 +156,8 @@ actual class ConsumerJson actual constructor(val topDir: String, val group: Grou
             return errs.add("'$ballotFilename' file does not exist")
         }
         return try {
-            fileSystemProvider.newInputStream(fileSystem.getPath(ballotFilename), StandardOpenOption.READ).use { inp ->
-                val json = Json.decodeFromStream<EncryptedBallotJson>(inp)
-                val eballot = json.import(group, errs)
-                if (errs.hasErrors()) Err(errs) else Ok(eballot!!)
-            }
+            val eballot = readEncryptedBallot(fileSystem.getPath(ballotFilename), errs)
+            if (errs.hasErrors()) Err(errs) else Ok(eballot!!)
         } catch (t: Throwable) {
             errs.add("Exception= ${t.message} ${t.stackTraceToString()}")
         }

egklib/src/jvmMain/kotlin/electionguard/publish/ConsumerProto.kt

@@ -8,16 +8,21 @@ import electionguard.ballot.*
 import electionguard.core.GroupContext
 import electionguard.core.fileReadBytes
 import electionguard.decrypt.DecryptingTrusteeIF
+import electionguard.json2.EncryptedBallotJson
+import electionguard.json2.import
 import electionguard.pep.BallotPep
 import electionguard.protoconvert.import
 import electionguard.util.ErrorMessages
 import io.github.oshai.kotlinlogging.KotlinLogging
+import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.decodeFromStream
 import pbandk.decodeFromByteBuffer
 import pbandk.decodeFromStream
 import java.io.*
 import java.nio.ByteBuffer
 import java.nio.file.Files
 import java.nio.file.Path
+import java.nio.file.StandardOpenOption
 import java.util.function.Predicate
 import java.util.stream.Stream
 
@@ -64,6 +69,13 @@ actual class ConsumerProto actual constructor(val topDir: String, val groupConte
         return groupContext.readDecryptionResult(protoPaths.decryptionResultPath())
     }
 
+    //////////////////////////////////////////////////////////////////////////
+
+    actual override fun hasEncryptedBallots(): Boolean {
+        val iter = iterateAllEncryptedBallots { true }
+        return iter.iterator().hasNext()
+    }
+
     actual override fun encryptingDevices(): List<String> {
         val topBallotPath = Path.of(protoPaths.encryptedBallotDir())
         if (!Files.exists(topBallotPath)) {
@@ -91,6 +103,20 @@ actual class ConsumerProto actual constructor(val topDir: String, val groupConte
         }
     }
 
+    actual override fun readEncryptedBallot(ballotDir: String, ballotId: String) : Result<EncryptedBallot, ErrorMessages> {
+        val errs = ErrorMessages("readEncryptedBallot ballotId=$ballotId from directory $ballotDir")
+        val ballotFilename = protoPaths.encryptedBallotPath(ballotDir, ballotId)
+        if (!Files.exists(Path.of(ballotFilename))) {
+            return errs.add("'$ballotFilename' file does not exist")
+        }
+        return try {
+            val eballot: EncryptedBallot? = readEncryptedBallot(ballotFilename, errs)
+            if (errs.hasErrors()) Err(errs) else Ok(eballot!!)
+        } catch (t: Throwable) {
+            errs.add("Exception= ${t.message} ${t.stackTraceToString()}")
+        }
+    }
+
     actual override fun iterateEncryptedBallots(
         device: String,
         filter: ((EncryptedBallot) -> Boolean)?
@@ -114,6 +140,14 @@ actual class ConsumerProto actual constructor(val topDir: String, val groupConte
         return Iterable { EncryptedBallotFileIterator(dirPath, filter) }
     }
 
+    actual override fun iterateAllEncryptedBallots(filter: ((EncryptedBallot) -> Boolean)?): Iterable<EncryptedBallot> {
+        val devices = encryptingDevices()
+        return Iterable { DeviceIterator(devices.iterator(), filter) }
+    }
+
+    //////////////////////////////////////////////////////////////////////////////////////////////
+
+
     actual override fun iterateDecryptedBallots(): Iterable<DecryptedTallyOrBallot> {
         // use batch (all in one protobuf) if it exists
         val batchedFileName: String = protoPaths.decryptedBatchPath()
@@ -129,16 +163,6 @@ actual class ConsumerProto actual constructor(val topDir: String, val groupConte
         return emptyList()
     }
 
-    actual override fun iterateAllEncryptedBallots(filter: ((EncryptedBallot) -> Boolean)?): Iterable<EncryptedBallot> {
-        val devices = encryptingDevices()
-        return Iterable { DeviceIterator(devices.iterator(), filter) }
-    }
-
-    actual override fun hasEncryptedBallots(): Boolean {
-        val iter = iterateAllEncryptedBallots { true }
-        return iter.iterator().hasNext()
-    }
-
     // plaintext ballots in given directory, with filter
     actual override fun iteratePlaintextBallots(
         ballotDir: String,
@@ -159,15 +183,7 @@ actual class ConsumerProto actual constructor(val topDir: String, val groupConte
         return groupContext.readTrustee(filename)
     }
 
-    actual override fun readEncryptedBallot(
-        ballotDir: String,
-        ballotId: String
-    ): Result<EncryptedBallot, ErrorMessages> {
-        val errs = ErrorMessages("readEncryptedBallot '$ballotDir/$ballotId'")
-        return errs.add("Not implemented yet")
-    }
-
-
+    // TODO havent defined PEP protos yet
     actual override fun iteratePepBallots(pepDir: String): Iterable<BallotPep> {
         throw RuntimeException("Not implemented yet")
     }
@@ -304,13 +320,9 @@ actual class ConsumerProto actual constructor(val topDir: String, val groupConte
         override fun computeNext() {
             while (true) {
                 if (ballotIds.hasNext()) {
-                    val ballotFile = protoPaths.encryptedBallotPath(device, ballotIds.next())
-                    var proto: electionguard.protogen.EncryptedBallot
-                    FileInputStream(ballotFile).use { inp ->
-                        proto = electionguard.protogen.EncryptedBallot.decodeFromStream(inp)
-                    }
+                    val ballotFile = protoPaths.encryptedBallotDevicePath(device, ballotIds.next())
                     val errs = ErrorMessages("EncryptedBallotChainIterator file='$ballotFile'")
-                    val result: EncryptedBallot? = proto.import(groupContext, errs)
+                    val result: EncryptedBallot? = readEncryptedBallot(ballotFile, errs)
                     if (errs.hasErrors()) {
                         logger.error { errs.toString() }
                         continue
@@ -375,12 +387,8 @@ actual class ConsumerProto actual constructor(val topDir: String, val groupConte
         override fun computeNext() {
             while (idx < pathList.size) {
                 val ballotFile = pathList[idx++]
-                var proto: electionguard.protogen.EncryptedBallot
-                FileInputStream(ballotFile.toString()).use { inp ->
-                    proto = electionguard.protogen.EncryptedBallot.decodeFromStream(inp)
-                }
                 val errs = ErrorMessages("EncryptedBallotFileIterator file='$ballotFile'")
-                val result: EncryptedBallot? = proto.import(groupContext, errs)
+                val result: EncryptedBallot? = readEncryptedBallot(ballotFile.toString(), errs)
                 if (errs.hasErrors()) {
                     logger.error { errs.toString() }
                     continue
@@ -396,6 +404,13 @@ actual class ConsumerProto actual constructor(val topDir: String, val groupConte
         }
     }
 
+    private fun readEncryptedBallot(ballotFile : String, errs: ErrorMessages): EncryptedBallot? {
+        FileInputStream(ballotFile).use { inp ->
+            val proto = electionguard.protogen.EncryptedBallot.decodeFromStream(inp)
+            return proto.import(groupContext, errs)
+        }
+    }
+
     private inner class DeviceIterator(
         val devices: Iterator<String>,
         private val filter: ((EncryptedBallot) -> Boolean)?,

egklib/src/jvmMain/kotlin/electionguard/publish/PublisherProto.kt

@@ -146,7 +146,7 @@ actual class PublisherProto actual constructor(topDir: String, createNew: Boolea
     inner class EncryptedBallotDeviceSink(val device: String) : EncryptedBallotSinkIF {
 
         override fun writeEncryptedBallot(ballot: EncryptedBallot) {
-            val ballotFile = protoPaths.encryptedBallotPath(device, ballot.ballotId)
+            val ballotFile = protoPaths.encryptedBallotDevicePath(device, ballot.ballotId)
             val ballotProto: pbandk.Message = ballot.publishProto()
             FileOutputStream(ballotFile).use { out -> ballotProto.encodeToStream(out) }
         }