From 43dbbbcde87e3ab736f93bf448d65f88b5c9ac27 Mon Sep 17 00:00:00 2001
From: Anton Firszov <antonfir@gmail.com>
Date: Wed, 10 Jun 2020 20:23:51 +0200
Subject: [PATCH] FixedCapacityPooledMemoryStreamTests

---
 .../IO/FixedCapacityPooledMemoryStream.cs     | 13 ++++++
 .../FixedCapacityPooledMemoryStreamTests.cs   | 42 +++++++++++++++++++
 2 files changed, 55 insertions(+)
 create mode 100644 tests/ImageSharp.Tests/IO/FixedCapacityPooledMemoryStreamTests.cs

diff --git a/src/ImageSharp/IO/FixedCapacityPooledMemoryStream.cs b/src/ImageSharp/IO/FixedCapacityPooledMemoryStream.cs
index 6953b0e10..74864d45e 100644
--- a/src/ImageSharp/IO/FixedCapacityPooledMemoryStream.cs
+++ b/src/ImageSharp/IO/FixedCapacityPooledMemoryStream.cs
@@ -1,6 +1,7 @@
 // Copyright (c) Six Labors.
 // Licensed under the Apache License, Version 2.0.
 
+using System;
 using System.Buffers;
 using System.IO;
 using SixLabors.ImageSharp.Memory;
@@ -29,6 +30,18 @@ namespace SixLabors.ImageSharp.IO
         /// <inheritdoc/>
         public override long Length { get; }
 
+        /// <inheritdoc/>
+        public override bool TryGetBuffer(out ArraySegment<byte> buffer)
+        {
+            if (this.isDisposed)
+            {
+                throw new ObjectDisposedException(this.GetType().Name);
+            }
+
+            buffer = new ArraySegment<byte>(this.buffer.Array, 0, this.buffer.Length());
+            return true;
+        }
+
         /// <inheritdoc/>
         protected override void Dispose(bool disposing)
         {
diff --git a/tests/ImageSharp.Tests/IO/FixedCapacityPooledMemoryStreamTests.cs b/tests/ImageSharp.Tests/IO/FixedCapacityPooledMemoryStreamTests.cs
new file mode 100644
index 000000000..0581a6ee2
--- /dev/null
+++ b/tests/ImageSharp.Tests/IO/FixedCapacityPooledMemoryStreamTests.cs
@@ -0,0 +1,42 @@
+// Copyright (c) Six Labors.
+// Licensed under the Apache License, Version 2.0.
+
+using System;
+using System.IO;
+using System.Linq;
+using SixLabors.ImageSharp.Memory;
+using SixLabors.ImageSharp.Tests.Memory;
+using Xunit;
+
+namespace SixLabors.ImageSharp.Tests.IO
+{
+    public class FixedCapacityPooledMemoryStreamTests
+    {
+        private readonly TestMemoryAllocator memoryAllocator = new TestMemoryAllocator();
+
+        [Theory]
+        [InlineData(1)]
+        [InlineData(512)]
+        public void RentsManagedBuffer(int length)
+        {
+            MemoryStream ms = this.memoryAllocator.AllocateFixedCapacityMemoryStream(length);
+            Assert.Equal(length, this.memoryAllocator.AllocationLog.Single().Length);
+            ms.Dispose();
+            Assert.Equal(1, this.memoryAllocator.ReturnLog.Count);
+        }
+
+        [Theory]
+        [InlineData(42)]
+        [InlineData(2999)]
+        public void UsesRentedBuffer(int length)
+        {
+            using MemoryStream ms = this.memoryAllocator.AllocateFixedCapacityMemoryStream(length);
+            ms.TryGetBuffer(out ArraySegment<byte> buffer);
+            byte[] array = buffer.Array;
+            Assert.Equal(array.GetHashCode(), this.memoryAllocator.AllocationLog.Single().HashCodeOfBuffer);
+
+            ms.Write(new byte[] { 123 });
+            Assert.Equal(123, array[0]);
+        }
+    }
+}