diff --git a/PowerKit.Tests/MemoryWriteStreamTests.cs b/PowerKit.Tests/MemoryWriteStreamTests.cs index e23bedb..5e978b6 100644 --- a/PowerKit.Tests/MemoryWriteStreamTests.cs +++ b/PowerKit.Tests/MemoryWriteStreamTests.cs @@ -1,3 +1,4 @@ +using System; using System.IO; using FluentAssertions; using PowerKit.Tests.Utils; @@ -29,4 +30,46 @@ public void MemoryWriteStream_Test() seekable.CanSeek.Should().BeTrue(); destination.ToArray().Should().Equal(data); } + + [Fact] + public void MemoryWriteStream_FlushTwice_Throws() + { + // Arrange + using var destination = new MemoryStream(); + using var seekable = new MemoryWriteStream(destination); + + seekable.Write([1, 2, 3]); + seekable.Flush(); + + // Act & Assert + seekable.Invoking(s => s.Flush()).Should().Throw(); + } + + [Fact] + public void MemoryWriteStream_WriteAfterFlush_Throws() + { + // Arrange + using var destination = new MemoryStream(); + using var seekable = new MemoryWriteStream(destination); + + seekable.Write([1, 2, 3]); + seekable.Flush(); + + // Act & Assert + seekable.Invoking(s => s.Write([4, 5, 6])).Should().Throw(); + } + + [Fact] + public void MemoryWriteStream_FlushThenDispose_DoesNotThrow() + { + // Arrange + using var destination = new MemoryStream(); + var seekable = new MemoryWriteStream(destination); + + seekable.Write([1, 2, 3]); + seekable.Flush(); + + // Act & Assert + seekable.Invoking(s => s.Dispose()).Should().NotThrow(); + } } diff --git a/PowerKit/MemoryWriteStream.cs b/PowerKit/MemoryWriteStream.cs index 6c01e56..5e34c69 100644 --- a/PowerKit/MemoryWriteStream.cs +++ b/PowerKit/MemoryWriteStream.cs @@ -5,16 +5,17 @@ namespace PowerKit; /// /// A wrapper that buffers all writes in memory and flushes them to the -/// underlying stream on . +/// underlying stream on . /// /// /// Writes go to an in-memory buffer and do not touch the underlying stream until -/// is called. This makes the wrapper always seekable and allows writes to +/// is called. This makes the wrapper always seekable and allows writes to /// be reordered freely before the final flush. /// public class MemoryWriteStream(Stream source) : Stream { private readonly MemoryStream _buffer = new(); + private bool _isFlushed; /// public override bool CanRead => false; @@ -35,12 +36,34 @@ public override long Position set => _buffer.Position = value; } - /// - public override void Flush() + private void Flush(bool throwIfAlreadyFlushed) { + if (_isFlushed) + { + if (throwIfAlreadyFlushed) + { + throw new InvalidOperationException( + "This stream has already been flushed and cannot be flushed again." + ); + } + + return; + } + _buffer.Position = 0; _buffer.CopyTo(source); source.Flush(); + _isFlushed = true; + } + + /// + public override void Flush() => Flush(true); + + /// + protected override void Dispose(bool disposing) + { + Flush(false); + base.Dispose(disposing); } /// @@ -54,6 +77,15 @@ public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException(); /// - public override void Write(byte[] buffer, int offset, int count) => + public override void Write(byte[] buffer, int offset, int count) + { + if (_isFlushed) + { + throw new InvalidOperationException( + "This stream has already been flushed and cannot be flushed again." + ); + } + _buffer.Write(buffer, offset, count); + } }