From be86b6c3324fffffd32386966b853565c3bcdd39 Mon Sep 17 00:00:00 2001
From: Jack Grigg <jack@z.cash>
Date: Wed, 31 Aug 2016 18:57:32 +1200
Subject: [PATCH] Add test for GetNoteDecryptor()

---
 src/gtest/test_keystore.cpp  | 12 ++++++++++++
 src/zcash/NoteEncryption.hpp |  1 +
 2 files changed, 13 insertions(+)

diff --git a/src/gtest/test_keystore.cpp b/src/gtest/test_keystore.cpp
index 8587ba49b..a5a952e9a 100644
--- a/src/gtest/test_keystore.cpp
+++ b/src/gtest/test_keystore.cpp
@@ -27,3 +27,15 @@ TEST(keystore_tests, store_and_retrieve_spending_key) {
     EXPECT_EQ(1, addrs.size());
     EXPECT_EQ(1, addrs.count(addr));
 }
+
+TEST(keystore_tests, store_and_retrieve_note_decryptor) {
+    CBasicKeyStore keyStore;
+    ZCNoteDecryption decOut;
+
+    auto sk = libzcash::SpendingKey::random();
+    auto addr = sk.address();
+    keyStore.AddSpendingKey(sk);
+
+    keyStore.GetNoteDecryptor(addr, decOut);
+    EXPECT_EQ(ZCNoteDecryption(sk.viewing_key()), decOut);
+}
diff --git a/src/zcash/NoteEncryption.hpp b/src/zcash/NoteEncryption.hpp
index 1cd1a9b27..bfeb80efa 100644
--- a/src/zcash/NoteEncryption.hpp
+++ b/src/zcash/NoteEncryption.hpp
@@ -71,6 +71,7 @@ public:
                       unsigned char nonce
                      ) const;
 
+    friend inline bool operator==(const NoteDecryption& a, const NoteDecryption& b) { return a.sk_enc == b.sk_enc && a.pk_enc == b.pk_enc; }
     friend inline bool operator<(const NoteDecryption& a, const NoteDecryption& b) { return a.pk_enc < b.pk_enc; }
 };