add missing test, simplify encrypted byte near limit calculation

1 files changed, 15 insertions, 3 deletions

diff --git a/tests/hazmat/primitives/test_aes.py b/tests/hazmat/primitives/test_aes.py
index 06e8adf8..f079ae4f 100644
--- a/tests/hazmat/primitives/test_aes.py
+++ b/tests/hazmat/primitives/test_aes.py
@@ -260,9 +260,7 @@ class TestAESModeGCM(object):
             modes.GCM(b"\x01" * 16),
             backend=backend
         ).encryptor()
-        # 16 bytes less than the encryption limit
-        near_limit_bytes = (2 ** 39 - 256 - 128) // 8
-        encryptor._bytes_processed = near_limit_bytes
+        encryptor._bytes_processed = modes.GCM._MAX_ENCRYPTED_BYTES - 16
         encryptor.update(b"0" * 16)
         assert (
             encryptor._bytes_processed == modes.GCM._MAX_ENCRYPTED_BYTES
@@ -270,6 +268,20 @@ class TestAESModeGCM(object):
         with pytest.raises(ValueError):
             encryptor.update(b"0")
 
+    def test_gcm_aad_limit(self, backend):
+        encryptor = base.Cipher(
+            algorithms.AES(b"\x00" * 16),
+            modes.GCM(b"\x01" * 16),
+            backend=backend
+        ).encryptor()
+        # 16 bytes less than the AAD limit
+        near_limit_bytes = (2 ** 64 - 128) // 8
+        encryptor._aad_bytes_processed = near_limit_bytes
+        encryptor.authenticate_additional_data(b"0" * 16)
+        assert encryptor._aad_bytes_processed == modes.GCM._MAX_AAD_BYTES
+        with pytest.raises(ValueError):
+            encryptor.authenticate_additional_data(b"0")
+
     def test_gcm_ciphertext_increments(self, backend):
         encryptor = base.Cipher(
             algorithms.AES(b"\x00" * 16),