From 90002459c40673ca702c7b89cd398eb6afd3547f Mon Sep 17 00:00:00 2001
From: Matthew Chen <matthew@signal.org>
Date: Mon, 6 Aug 2018 13:02:01 -0400
Subject: [PATCH] Add unit tests around proto wrappers.

---
 Signal/test/SSKTests/SSKEnvelopeTest.swift    | 49 +++++++++++++++++++
 .../test/call/PeerConnectionClientTest.swift  | 14 ++++--
 2 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/Signal/test/SSKTests/SSKEnvelopeTest.swift b/Signal/test/SSKTests/SSKEnvelopeTest.swift
index 056a64b29..58206520a 100644
--- a/Signal/test/SSKTests/SSKEnvelopeTest.swift
+++ b/Signal/test/SSKTests/SSKEnvelopeTest.swift
@@ -65,4 +65,53 @@ class SSKProtoEnvelopeTest: XCTestCase {
             }
         }
     }
+
+    func testParse_buildVsRequired() {
+        let builder = SSKProtoEnvelope.SSKProtoEnvelopeBuilder()
+
+        XCTAssertThrowsError(try builder.build()) { (error) -> Void in
+            switch error {
+            case SSKProtoError.invalidProtobuf:
+                break
+            default:
+                XCTFail("unexpected error: \(error)")
+            }
+        }
+    }
+
+    func testParse_roundtrip() {
+        let builder = SSKProtoEnvelope.SSKProtoEnvelopeBuilder()
+
+        let phonyContent = "phony data".data(using: .utf8)!
+
+        builder.setType(SSKProtoEnvelope.SSKProtoEnvelopeType.prekeyBundle)
+        builder.setTimestamp(123)
+        builder.setSource("+13213214321")
+        builder.setSourceDevice(1)
+        builder.setContent(phonyContent)
+
+        var envelopeData: Data
+        do {
+            envelopeData = try builder.buildSerializedData()
+        } catch {
+            XCTFail("Couldn't serialize data.")
+            return
+        }
+
+        var envelope: SSKProtoEnvelope
+        do {
+            envelope = try SSKProtoEnvelope.parseData(envelopeData)
+        } catch {
+            XCTFail("Couldn't serialize data.")
+            return
+        }
+
+        XCTAssertEqual(envelope.type, SSKProtoEnvelope.SSKProtoEnvelopeType.prekeyBundle)
+        XCTAssertEqual(envelope.timestamp, 123)
+        XCTAssertEqual(envelope.source, "+13213214321")
+        XCTAssertEqual(envelope.sourceDevice, 1)
+        XCTAssertTrue(envelope.hasContent)
+        XCTAssertEqual(envelope.content, phonyContent)
+        XCTAssertFalse(envelope.hasLegacyMessage)
+    }
 }
diff --git a/Signal/test/call/PeerConnectionClientTest.swift b/Signal/test/call/PeerConnectionClientTest.swift
index a2f5d44ff..20e1ded1b 100644
--- a/Signal/test/call/PeerConnectionClientTest.swift
+++ b/Signal/test/call/PeerConnectionClientTest.swift
@@ -115,8 +115,16 @@ class PeerConnectionClientTest: XCTestCase {
     func testDataChannelMessage() {
         XCTAssertEqual(0, clientDelegate.dataChannelMessages.count)
 
-        let hangup = DataChannelMessage.forHangup(callId: 123)
-        let hangupBuffer = RTCDataBuffer(data: hangup.asData(), isBinary: false)
+        var hangupData: Data
+        do {
+            let hangupBuilder = WebRTCProtoHangup.WebRTCProtoHangupBuilder()
+            hangupBuilder.setId(123)
+            hangupData = try hangupBuilder.buildSerializedData()
+        } catch {
+            XCTFail("Couldn't build proto.")
+            return
+        }
+        let hangupBuffer = RTCDataBuffer(data: hangupData, isBinary: false)
         client.dataChannel(dataChannel, didReceiveMessageWith: hangupBuffer)
 
         waitForPeerConnectionClient()
@@ -124,7 +132,7 @@ class PeerConnectionClientTest: XCTestCase {
         XCTAssertEqual(1, clientDelegate.dataChannelMessages.count)
 
         let dataChannelMessageProto = clientDelegate.dataChannelMessages[0]
-        XCTAssert(dataChannelMessageProto.hasHangup())
+        XCTAssert(dataChannelMessageProto.hasHangup)
 
         let hangupProto = dataChannelMessageProto.hangup!
         XCTAssertEqual(123, hangupProto.id)