From c8c232f196e0c19790dc8bb5b22f4da7159fa8d1 Mon Sep 17 00:00:00 2001
From: Mark Hansen <markhansen@google.com>
Date: Sun, 22 Sep 2024 17:14:08 -0700
Subject: [PATCH] Add a test for writing a single byte, including beyond the
 end of the bounds.

This test doesn't yet pass for some OutputTypes; we'll fix that in the next CL.

PiperOrigin-RevId: 677565048
---
 .../protobuf/CodedOutputStreamTest.java       | 37 +++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/java/core/src/test/java/com/google/protobuf/CodedOutputStreamTest.java b/java/core/src/test/java/com/google/protobuf/CodedOutputStreamTest.java
index 19f3b5d2ca..d6ff12df85 100644
--- a/java/core/src/test/java/com/google/protobuf/CodedOutputStreamTest.java
+++ b/java/core/src/test/java/com/google/protobuf/CodedOutputStreamTest.java
@@ -209,6 +209,12 @@ public class CodedOutputStreamTest {
     };
 
     abstract Coder newCoder(int size);
+
+    /** Whether we can call CodedOutputStream.spaceLeft(). */
+    boolean supportsSpaceLeft() {
+      // STREAM doesn't know how much space is left.
+      return this != OutputType.STREAM;
+    }
   }
 
   /** Checks that invariants are maintained for varint round trip input and output. */
@@ -526,6 +532,37 @@ public class CodedOutputStreamTest {
     }
   }
 
+  @Test
+  public void testWriteByte() throws Exception {
+    Coder coder = outputType.newCoder(5);
+    // Write 5 bytes
+    coder.stream().write((byte) 1);
+    coder.stream().write((byte) 1);
+    coder.stream().write((byte) 1);
+    coder.stream().write((byte) 1);
+    coder.stream().write((byte) 1);
+    coder.stream().flush();
+    byte[] rawBytes = coder.toByteArray();
+    assertThat(rawBytes).isEqualTo(new byte[] {1, 1, 1, 1, 1});
+    if (outputType.supportsSpaceLeft()) {
+      assertThat(coder.stream().spaceLeft()).isEqualTo(0);
+    }
+
+    // Going beyond bounds should throw. Except if we're streaming, where buffering masks the
+    // failure.
+    if (outputType == OutputType.STREAM) {
+      return;
+    }
+    assertThrows(OutOfSpaceException.class, () -> coder.stream().write((byte) 1));
+    // For some OutputTypes, this test doesn't pass yet.
+    if (outputType.supportsSpaceLeft()
+        && outputType != OutputType.ARRAY
+        && outputType != OutputType.NIO_HEAP
+        && outputType != OutputType.NIO_HEAP_WITH_INITIAL_OFFSET) {
+      assertThat(coder.stream().spaceLeft()).isEqualTo(0);
+    }
+  }
+
   @Test
   public void testWriteByteBuffer() throws Exception {
     byte[] value = "abcde".getBytes(Internal.UTF_8);