From c8b7e8713543d7110766697b0ca2fbc82f33f084 Mon Sep 17 00:00:00 2001
From: wussler <aron@wussler.it>
Date: Thu, 5 Mar 2020 22:16:38 +0100
Subject: [PATCH] Fix test encryption test (#37)

* Fix TestMultipleKeyMessageEncryption

* Fix TestMultipleKeyMessageEncryption - simplify code

* Add diffs to changelog

* Simplify test code further
---
 CHANGELOG.md           |  4 ++++
 crypto/message_test.go | 24 +++++++++++++++---------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 21c91c6..c8a6a27 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+### Fixed
+- Fixed test `TestMultipleKeyMessageEncryption`
+
 ## [2.0.0] - 2020-01-06
 Since the open-sourcing of the library in May the API has been updated, listening to internal and
 external feedback, in order to have a flexible library, that can be used in a simple settings,
diff --git a/crypto/message_test.go b/crypto/message_test.go
index 5e52355..b38c171 100644
--- a/crypto/message_test.go
+++ b/crypto/message_test.go
@@ -3,7 +3,6 @@ package crypto
 import (
 	"bytes"
 	"encoding/base64"
-	"io"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -153,19 +152,26 @@ func TestMultipleKeyMessageEncryption(t *testing.T) {
 		t.Fatal("Expected no error when encrypting, got:", err)
 	}
 
-	numKeyPackets := 0
+	// Test that ciphertext data contains three Encrypted Key Packets (tag 1)
+	// followed by a single symmetrically encrypted data packet (tag 18)
+	var p packet.Packet
 	packets := packet.NewReader(bytes.NewReader(ciphertext.Data))
-	for {
-		var p packet.Packet
-		if p, err = packets.Next(); err == io.EOF {
-			break
+	for i := 0; i < 3; i++ {
+		if p, err = packets.Next(); err != nil {
+			t.Fatal(err.Error())
 		}
-		if _, ok := p.(*packet.EncryptedKey); ok {
-			numKeyPackets++
+		if _, ok := p.(*packet.EncryptedKey); !ok {
+			t.Fatalf("Expected Encrypted Key packet, got %T", p)
 		}
 	}
-	assert.Exactly(t, 3, numKeyPackets)
+	if p, err = packets.Next(); err != nil {
+		t.Fatal(err.Error())
+	}
+	if _, ok := p.(*packet.SymmetricallyEncrypted); !ok {
+		t.Fatalf("Expected Symmetrically Encrypted Data packet, got %T", p)
+	}
 
+	// Decrypt message and verify correctness
 	decrypted, err := keyRingTestPrivate.Decrypt(ciphertext, keyRingTestPublic, GetUnixTime())
 	if err != nil {
 		t.Fatal("Expected no error when decrypting, got:", err)