1 files changed, 58 insertions, 4 deletions

diff --git a/src/DotNetOpenAuth.Test/TestBase.cs b/src/DotNetOpenAuth.Test/TestBase.cs
index 92adafa..4758b7d 100644
--- a/src/DotNetOpenAuth.Test/TestBase.cs
+++ b/src/DotNetOpenAuth.Test/TestBase.cs
@@ -7,7 +7,12 @@
 namespace DotNetOpenAuth.Test {
 	using System;
 	using System.IO;
+	using System.Net;
+	using System.Net.Http;
+	using System.Net.Http.Headers;
 	using System.Reflection;
+	using System.Threading;
+	using System.Threading.Tasks;
 	using System.Web;
 	using DotNetOpenAuth.Messaging.Reflection;
 	using DotNetOpenAuth.OAuth.Messages;
@@ -16,6 +21,8 @@ namespace DotNetOpenAuth.Test {
 	using log4net;
 	using NUnit.Framework;
 
+	using log4net.Config;
+
 	/// <summary>
 	/// The base class that all test classes inherit from.
 	/// </summary>
@@ -48,14 +55,17 @@ namespace DotNetOpenAuth.Test {
 			get { return this.messageDescriptions; }
 		}
 
+		internal MockingHostFactories HostFactories;
+
 		/// <summary>
 		/// The TestInitialize method for the test cases.
 		/// </summary>
 		[SetUp]
 		public virtual void SetUp() {
-			log4net.Config.XmlConfigurator.Configure(Assembly.GetExecutingAssembly().GetManifestResourceStream("DotNetOpenAuth.Test.Logging.config"));
+			XmlConfigurator.Configure(Assembly.GetExecutingAssembly().GetManifestResourceStream("DotNetOpenAuth.Test.Logging.config"));
 			MessageBase.LowSecurityMode = true;
 			this.messageDescriptions = new MessageDescriptionCollection();
+			this.HostFactories = new MockingHostFactories();
 			SetMockHttpContext();
 		}
 
@@ -64,10 +74,10 @@ namespace DotNetOpenAuth.Test {
 		/// </summary>
 		[TearDown]
 		public virtual void Cleanup() {
-			log4net.LogManager.Shutdown();
+			LogManager.Shutdown();
 		}
 
-		internal static Stats MeasurePerformance(Action action, float maximumAllowedUnitTime, int samples = 10, int iterations = 100, string name = null) {
+		internal static Stats MeasurePerformance(Func<Task> action, float maximumAllowedUnitTime, int samples = 10, int iterations = 100, string name = null) {
 			if (!PerformanceTestUtilities.IsOptimized(typeof(OpenIdRelyingParty).Assembly)) {
 				Assert.Inconclusive("Unoptimized code.");
 			}
@@ -75,7 +85,7 @@ namespace DotNetOpenAuth.Test {
 			var timer = new MultiSampleCodeTimer(samples, iterations);
 			Stats stats;
 			using (new HighPerformance()) {
-				stats = timer.Measure(name ?? TestContext.CurrentContext.Test.FullName, action);
+				stats = timer.Measure(name ?? TestContext.CurrentContext.Test.FullName, () => action().Wait());
 			}
 
 			stats.AdjustForScale(PerformanceTestUtilities.Baseline.Median);
@@ -103,5 +113,49 @@ namespace DotNetOpenAuth.Test {
 				new HttpRequest("mock", "http://mock", "mock"),
 				new HttpResponse(new StringWriter()));
 		}
+
+		protected internal Handler Handle(string uri) {
+			return new Handler(this, new Uri(uri));
+		}
+
+		protected internal Handler Handle(Uri uri) {
+			return new Handler(this, uri);
+		}
+
+		protected internal struct Handler {
+			private TestBase test;
+
+			internal Handler(TestBase test, Uri uri)
+				: this() {
+				this.test = test;
+				this.Uri = uri;
+			}
+
+			internal Uri Uri { get; private set; }
+
+			internal Func<HttpRequestMessage, Task<HttpResponseMessage>> MessageHandler { get; private set; }
+
+			internal void By(Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> handler) {
+				this.test.HostFactories.Handlers[this.Uri] = req => handler(req, CancellationToken.None);
+			}
+
+			internal void By(Func<HttpRequestMessage, Task<HttpResponseMessage>> handler) {
+				this.test.HostFactories.Handlers[this.Uri] = handler;
+			}
+
+			internal void By(Func<HttpRequestMessage, HttpResponseMessage> handler) {
+				this.By(req => Task.FromResult(handler(req)));
+			}
+
+			internal void By(string responseContent, string contentType, HttpStatusCode statusCode = HttpStatusCode.OK) {
+				this.By(
+				   req => {
+					   var response = new HttpResponseMessage(statusCode);
+					   response.Content = new StringContent(responseContent);
+					   response.Content.Headers.ContentType = new MediaTypeHeaderValue(contentType);
+					   return response;
+				   });
+			}
+		}
 	}
 }