From fd4870dd3900f9a65858141086aa0f0b77046797 Mon Sep 17 00:00:00 2001
From: David Wambugu <dawambug@microsoft.com>
Date: Tue, 10 Sep 2024 12:48:17 +0300
Subject: [PATCH] Update code from pr review.

---
 .../MessageWriterConcurrencyTests.cs          | 51 +++----------------
 1 file changed, 7 insertions(+), 44 deletions(-)

diff --git a/test/FunctionalTests/Microsoft.OData.Core.Tests/MessageWriterConcurrencyTests.cs b/test/FunctionalTests/Microsoft.OData.Core.Tests/MessageWriterConcurrencyTests.cs
index fe1728112d..35f842ab54 100644
--- a/test/FunctionalTests/Microsoft.OData.Core.Tests/MessageWriterConcurrencyTests.cs
+++ b/test/FunctionalTests/Microsoft.OData.Core.Tests/MessageWriterConcurrencyTests.cs
@@ -4,13 +4,13 @@
 // </copyright>
 //---------------------------------------------------------------------
 
-using System.Collections.Generic;
 using System.IO;
 using System.Threading.Tasks;
 using System;
 using Xunit;
 using Microsoft.Extensions.DependencyInjection;
 using System.Linq;
+using Microsoft.OData.Tests;
 
 namespace Microsoft.OData.Core.Tests
 {
@@ -52,7 +52,12 @@ private static async Task<string> WritePayloadAsync(string content, IServiceProv
         {
             using Stream outputStream = new MemoryStream();
 
-            var message = new ODataMessage(outputStream, serviceProvider);
+            var message = new InMemoryMessage
+            {
+                Stream = outputStream,
+                ServiceProvider = serviceProvider
+            };
+
             var responseMessage = new ODataResponseMessage(message, writing: true, enableMessageStreamDisposal: true, maxMessageSize: -1);
             await using ODataMessageWriter writer = new ODataMessageWriter(responseMessage);
 
@@ -69,47 +74,5 @@ private static async Task<string> WritePayloadAsync(string content, IServiceProv
             await writer.DisposeAsync();
             return written;
         }
-
-
-        class ODataMessage : IODataResponseMessage, IODataResponseMessageAsync, IServiceCollectionProvider
-        {
-            private Dictionary<string, string> _headers = new();
-            private Stream _outputStream;
-            public ODataMessage(Stream outputStream, IServiceProvider serviceProvider)
-            {
-                this.ServiceProvider = serviceProvider;
-                _outputStream = outputStream;
-            }
-            public IEnumerable<KeyValuePair<string, string>> Headers => _headers;
-
-            public int StatusCode { get => throw new NotImplementedException(); set => throw new NotImplementedException(); }
-
-            public IServiceProvider ServiceProvider { get; private set; }
-
-            public string GetHeader(string headerName)
-            {
-                if (_headers.TryGetValue(headerName, out var value))
-                {
-                    return value;
-                }
-
-                return null;
-            }
-
-            public Stream GetStream()
-            {
-                return _outputStream;
-            }
-
-            public Task<Stream> GetStreamAsync()
-            {
-                return Task.FromResult(_outputStream);
-            }
-
-            public void SetHeader(string headerName, string headerValue)
-            {
-                _headers[headerName] = headerValue;
-            }
-        }
     }
 }