Support product checking on virtual cluster

stevejgordon · stevejgordon · commit f5108094c8a0 · 2021-06-15T15:16:56.000+01:00
diff --git a/src/Elasticsearch.Net.VirtualizedCluster/Audit/Auditor.cs b/src/Elasticsearch.Net.VirtualizedCluster/Audit/Auditor.cs
@@ -53,16 +53,17 @@ public void ChangeTime(Func<DateTime, DateTime> selector)
 		public async Task<Auditor> TraceStartup(ClientCall callTrace = null)
 		{
 			//synchronous code path
-			_cluster = _cluster ?? Cluster();
+			_cluster ??= Cluster();
 			if (!StartedUp) AssertPoolBeforeStartup?.Invoke(_cluster.ConnectionPool);
 			AssertPoolBeforeCall?.Invoke(_cluster.ConnectionPool);
+			// ReSharper disable once MethodHasAsyncOverload
 			Response = _cluster.ClientCall(callTrace?.RequestOverrides);
 			AuditTrail = Response.ApiCall.AuditTrail;
 			if (!StartedUp) AssertPoolAfterStartup?.Invoke(_cluster.ConnectionPool);
 			AssertPoolAfterCall?.Invoke(_cluster.ConnectionPool);
 
 			//async code path
-			_clusterAsync = _clusterAsync ?? Cluster();
+			_clusterAsync ??= Cluster();
 			if (!StartedUp) AssertPoolBeforeStartup?.Invoke(_clusterAsync.ConnectionPool);
 			AssertPoolBeforeCall?.Invoke(_clusterAsync.ConnectionPool);
 			ResponseAsync = await _clusterAsync.ClientCallAsync(callTrace?.RequestOverrides).ConfigureAwait(false);
@@ -71,19 +72,25 @@ public async Task<Auditor> TraceStartup(ClientCall callTrace = null)
 			AssertPoolAfterCall?.Invoke(_clusterAsync.ConnectionPool);
 			return new Auditor(_cluster, _clusterAsync);
 		}
-
+		
 		public async Task<Auditor> TraceCall(ClientCall callTrace, int nthCall = 0)
 		{
 			await TraceStartup(callTrace).ConfigureAwait(false);
-			return AssertAuditTrails(callTrace, nthCall);
+			return AssertAuditTrails(callTrace, nthCall, true);
+		}
+
+		public async Task<Auditor> TraceCall(bool skipProductCheck, ClientCall callTrace, int nthCall = 0)
+		{
+			await TraceStartup(callTrace).ConfigureAwait(false);
+			return AssertAuditTrails(callTrace, nthCall, skipProductCheck);
 		}
 
 #pragma warning disable 1998 // Async method lacks 'await' operators and will run synchronously
 		private async Task TraceException<TException>(ClientCall callTrace, Action<TException> assert)
 #pragma warning restore 1998 // Async method lacks 'await' operators and will run synchronously
 			where TException : ElasticsearchClientException
 		{
-			_cluster = _cluster ?? Cluster();
+			_cluster ??= Cluster();
 			_cluster.ClientThrows(true);
 			AssertPoolBeforeCall?.Invoke(_cluster.ConnectionPool);
 
@@ -94,7 +101,7 @@ private async Task TraceException<TException>(ClientCall callTrace, Action<TExce
 			AuditTrail = exception.AuditTrail;
 			AssertPoolAfterCall?.Invoke(_cluster.ConnectionPool);
 
-			_clusterAsync = _clusterAsync ?? Cluster();
+			_clusterAsync ??= Cluster();
 			_clusterAsync.ClientThrows(true);
 			Func<Task> callAsync = async () => ResponseAsync = await _clusterAsync.ClientCallAsync(callTrace?.RequestOverrides).ConfigureAwait(false);
 			exception = await TryCallAsync(callAsync, assert).ConfigureAwait(false);
@@ -121,7 +128,7 @@ public async Task<Auditor> TraceUnexpectedElasticsearchException(ClientCall call
 		public async Task<Auditor> TraceElasticsearchExceptionOnResponse(ClientCall callTrace, Action<ElasticsearchClientException> assert)
 #pragma warning restore 1998
 		{
-			_cluster = _cluster ?? Cluster();
+			_cluster ??= Cluster();
 			_cluster.ClientThrows(false);
 			AssertPoolBeforeCall?.Invoke(_cluster.ConnectionPool);
 
@@ -130,14 +137,15 @@ public async Task<Auditor> TraceElasticsearchExceptionOnResponse(ClientCall call
 
 			if (Response.ApiCall.Success) throw new Exception("Expected call to not be valid");
 
-			var exception = Response.ApiCall.OriginalException as ElasticsearchClientException;
-			if (exception == null) throw new Exception("OriginalException on response is not expected ElasticsearchClientException");
+			if (Response.ApiCall.OriginalException is not ElasticsearchClientException exception)
+				throw new Exception("OriginalException on response is not expected ElasticsearchClientException");
+
 			assert(exception);
 
 			AuditTrail = exception.AuditTrail;
 			AssertPoolAfterCall?.Invoke(_cluster.ConnectionPool);
 
-			_clusterAsync = _clusterAsync ?? Cluster();
+			_clusterAsync ??= Cluster();
 			_clusterAsync.ClientThrows(false);
 			Func<Task> callAsync = async () => { ResponseAsync = await _clusterAsync.ClientCallAsync(callTrace?.RequestOverrides).ConfigureAwait(false); };
 			await callAsync().ConfigureAwait(false);
@@ -157,30 +165,38 @@ public async Task<Auditor> TraceElasticsearchExceptionOnResponse(ClientCall call
 		public async Task<Auditor> TraceUnexpectedException(ClientCall callTrace, Action<UnexpectedElasticsearchClientException> assert)
 #pragma warning restore 1998
 		{
-			_cluster = _cluster ?? Cluster();
+			_cluster ??= Cluster();
 			AssertPoolBeforeCall?.Invoke(_cluster.ConnectionPool);
 
-			Action call = () => Response = _cluster.ClientCall(callTrace?.RequestOverrides);
-			var exception = TryCall(call, assert);
+			
+			var exception = TryCall(Call, assert);
 			assert(exception);
 
 			AuditTrail = exception.AuditTrail;
 			AssertPoolAfterCall?.Invoke(_cluster.ConnectionPool);
 
-			_clusterAsync = _clusterAsync ?? Cluster();
-			Func<Task> callAsync = async () => ResponseAsync = await _clusterAsync.ClientCallAsync(callTrace?.RequestOverrides).ConfigureAwait(false);
-			exception = await TryCallAsync(callAsync, assert).ConfigureAwait(false);
+			_clusterAsync ??= Cluster();
+			
+			exception = await TryCallAsync(CallAsync, assert).ConfigureAwait(false);
 			assert(exception);
 
 			AsyncAuditTrail = exception.AuditTrail;
 			AssertPoolAfterCall?.Invoke(_clusterAsync.ConnectionPool);
+
 			return new Auditor(_cluster, _clusterAsync);
+
+			void Call() => Response = _cluster.ClientCall(callTrace?.RequestOverrides);
+			async Task CallAsync() => ResponseAsync = await _clusterAsync.ClientCallAsync(callTrace?.RequestOverrides).ConfigureAwait(false);
 		}
 
-		private Auditor AssertAuditTrails(ClientCall callTrace, int nthCall)
+		private Auditor AssertAuditTrails(ClientCall callTrace, int nthCall, bool skipProductCheck)
 		{
 			var nl = Environment.NewLine;
-			if (AuditTrail.Count != AsyncAuditTrail.Count)
+
+			if (skipProductCheck)
+				AuditTrail.RemoveAll(a => a.Event is AuditEvent.ProductCheckOnStartup or AuditEvent.ProductCheckSuccess or AuditEvent.ProductCheckFailure);
+
+			if (AuditTrail.Count(Predicate) != AsyncAuditTrail.Count(Predicate))
 				throw new Exception($"{nthCall} has a mismatch between sync and async. {nl}async:{AuditTrail}{nl}sync:{AsyncAuditTrail}");
 
 			AssertTrailOnResponse(callTrace, AuditTrail, true, nthCall);
@@ -192,8 +208,13 @@ private Auditor AssertAuditTrails(ClientCall callTrace, int nthCall)
 			callTrace?.AssertPoolAfterCall?.Invoke(_cluster.ConnectionPool);
 			callTrace?.AssertPoolAfterCall?.Invoke(_clusterAsync.ConnectionPool);
 			return new Auditor(_cluster, _clusterAsync);
-		}
 
+			// These happen one time only so should not be counted when comparing equality of audit trails
+			static bool Predicate(Net.Audit auditEvent) =>
+				auditEvent.Event != AuditEvent.ProductCheckOnStartup &&
+				auditEvent.Event != AuditEvent.ProductCheckFailure &&
+				auditEvent.Event != AuditEvent.ProductCheckSuccess;
+		}
 
 		public void VisualizeCalls(int numberOfCalls)
 		{
@@ -218,14 +239,24 @@ private static string AuditTrailToString(List<Elasticsearch.Net.Audit> auditTrai
 			return actualAuditTrail;
 		}
 
+
 		public async Task<Auditor> TraceCalls(params ClientCall[] audits)
 		{
 			var auditor = this;
-			foreach (var a in audits.Select((a, i) => new { a, i })) auditor = await auditor.TraceCall(a.a, a.i).ConfigureAwait(false);
+			foreach (var a in audits.Select((a, i) => new { a, i }))
+				auditor = await auditor.TraceCall(a.a, a.i).ConfigureAwait(false);
+			return auditor;
+		}
+
+		public async Task<Auditor> TraceCalls(bool skipProductCheck, params ClientCall[] audits)
+		{
+			var auditor = this;
+			foreach (var a in audits.Select((a, i) => new { a, i }))
+				auditor = await auditor.TraceCall(skipProductCheck, a.a, a.i).ConfigureAwait(false);
 			return auditor;
 		}
 
-		private static void AssertTrailOnResponse(ClientCall callTrace, List<Elasticsearch.Net.Audit> auditTrail, bool sync, int nthCall)
+		private static void AssertTrailOnResponse(ClientCall callTrace, IReadOnlyCollection<Net.Audit> auditTrail, bool sync, int nthCall)
 		{
 			var typeOfTrail = (sync ? "synchronous" : "asynchronous") + " audit trail";
 			var nthClientCall = (nthCall + 1).ToOrdinal();
diff --git a/src/Elasticsearch.Net.VirtualizedCluster/Rules/ProductCheckRule.cs b/src/Elasticsearch.Net.VirtualizedCluster/Rules/ProductCheckRule.cs
@@ -0,0 +1,33 @@
+﻿// Licensed to Elasticsearch B.V under one or more agreements.
+// Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+// See the LICENSE file in the project root for more information
+
+using System;
+
+namespace Elasticsearch.Net.VirtualizedCluster.Rules
+{
+	public class ProductCheckRule : RuleBase<ProductCheckRule>
+	{
+		private IRule Self => this;
+
+		public ProductCheckRule Fails(RuleOption<TimesHelper.AllTimes, int> times, RuleOption<Exception, int> errorState = null)
+		{
+			Self.Times = times;
+			Self.Succeeds = false;
+			Self.Return = errorState;
+			return this;
+		}
+
+		public ProductCheckRule Succeeds(RuleOption<TimesHelper.AllTimes, int> times, int? validResponseCode = 200)
+		{
+			Self.Times = times;
+			Self.Succeeds = true;
+			Self.Return = validResponseCode;
+			return this;
+		}
+
+		public ProductCheckRule SucceedAlways(int? validResponseCode = 200) => Succeeds(TimesHelper.Always, validResponseCode);
+
+		public ProductCheckRule FailAlways(RuleOption<Exception, int> errorState = null) => Fails(TimesHelper.Always, errorState);
+	}
+}
diff --git a/src/Elasticsearch.Net.VirtualizedCluster/VirtualCluster.cs b/src/Elasticsearch.Net.VirtualizedCluster/VirtualCluster.cs
@@ -15,15 +15,23 @@ public class VirtualCluster
 	{
 		private readonly List<Node> _nodes;
 
-		public VirtualCluster(IEnumerable<Node> nodes) => _nodes = nodes.ToList();
+		public VirtualCluster(IEnumerable<Node> nodes, bool productCheckSucceeds = true)
+		{
+			_nodes = nodes.ToList();
+
+			if (productCheckSucceeds)
+				ProductCheckRules.Add(new ProductCheckRule().SucceedAlways());
+		}
 
-		public List<IClientCallRule> ClientCallRules { get; } = new List<IClientCallRule>();
-		public TestableDateTimeProvider DateTimeProvider { get; } = new TestableDateTimeProvider();
+		public List<IClientCallRule> ClientCallRules { get; } = new();
+		public TestableDateTimeProvider DateTimeProvider { get; } = new();
 
 		public IReadOnlyList<Node> Nodes => _nodes;
-		public List<IRule> PingingRules { get; } = new List<IRule>();
 
-		public List<ISniffRule> SniffingRules { get; } = new List<ISniffRule>();
+		public List<IRule> PingingRules { get; } = new();
+		public List<ISniffRule> SniffingRules { get; } = new();
+		public List<IRule> ProductCheckRules { get; } = new();
+
 		internal string PublishAddressOverride { get; private set; }
 
 		internal bool SniffShouldReturnFqnd { get; private set; }
@@ -85,6 +93,12 @@ public VirtualCluster Sniff(Func<SniffRule, ISniffRule> selector)
 			SniffingRules.Add(selector(new SniffRule()));
 			return this;
 		}
+		
+		public VirtualCluster ProductCheck(Func<ProductCheckRule, IRule> selector)
+		{
+			ProductCheckRules.Add(selector(new ProductCheckRule()));
+			return this;
+		}
 
 		public VirtualCluster ClientCalls(Func<ClientCallRule, IClientCallRule> selector)
 		{
diff --git a/src/Elasticsearch.Net.VirtualizedCluster/VirtualClusterConnection.cs b/src/Elasticsearch.Net.VirtualizedCluster/VirtualClusterConnection.cs
@@ -6,16 +6,16 @@
 using System.Collections.Generic;
 using System.IO;
 using System.Linq;
-#if DOTNETCORE
-using TheException = System.Net.Http.HttpRequestException;
-#else
-using TheException = System.Net.WebException;
-#endif
 using System.Threading;
 using System.Threading.Tasks;
 using Elasticsearch.Net.VirtualizedCluster.MockResponses;
 using Elasticsearch.Net.VirtualizedCluster.Providers;
 using Elasticsearch.Net.VirtualizedCluster.Rules;
+#if DOTNETCORE
+using TheException = System.Net.Http.HttpRequestException;
+#else
+using TheException = System.Net.WebException;
+#endif
 
 namespace Elasticsearch.Net.VirtualizedCluster
 {
@@ -31,7 +31,7 @@ namespace Elasticsearch.Net.VirtualizedCluster
 	/// </summary>
 	public class VirtualClusterConnection : InMemoryConnection
 	{
-		private static readonly object Lock = new object();
+		private static readonly object Lock = new();
 
 		private static byte[] _defaultResponseBytes;
 
@@ -101,6 +101,9 @@ public bool IsPingRequest(RequestData requestData) =>
 			requestData.Method == HttpMethod.HEAD &&
 			(requestData.PathAndQuery == string.Empty || requestData.PathAndQuery.StartsWith("?"));
 
+		public bool IsProductCheckRequest(RequestData requestData) =>
+			requestData.Uri.AbsolutePath.Equals("/", StringComparison.Ordinal) && requestData.Method == HttpMethod.GET;
+
 		public override TResponse Request<TResponse>(RequestData requestData)
 		{
 			if (!_calls.ContainsKey(requestData.Uri.Port))
@@ -117,8 +120,8 @@ public override TResponse Request<TResponse>(RequestData requestData)
 						nameof(VirtualCluster.Sniff),
 						_cluster.SniffingRules,
 						requestData.RequestTimeout,
-						(r) => UpdateCluster(r.NewClusterState),
-						(r) => SniffResponseBytes.Create(_cluster.Nodes, _cluster.ElasticsearchVersion,_cluster.PublishAddressOverride, _cluster.SniffShouldReturnFqnd)
+						r => UpdateCluster(r.NewClusterState),
+						_ => SniffResponseBytes.Create(_cluster.Nodes, _cluster.ElasticsearchVersion,_cluster.PublishAddressOverride, _cluster.SniffShouldReturnFqnd)
 					);
 				}
 				if (IsPingRequest(requestData))
@@ -129,8 +132,20 @@ public override TResponse Request<TResponse>(RequestData requestData)
 						nameof(VirtualCluster.Ping),
 						_cluster.PingingRules,
 						requestData.PingTimeout,
-						(r) => { },
-						(r) => null //HEAD request
+						_ => { },
+						_ => null //HEAD request
+					);
+				}
+				if (IsProductCheckRequest(requestData))
+				{
+					_ = Interlocked.Increment(ref state.ProductChecked);
+					return HandleRules<TResponse, IRule>(
+						requestData,
+						nameof(VirtualCluster.ProductCheck),
+						_cluster.ProductCheckRules,
+						requestData.RequestTimeout,
+						_ => { },
+						_ => ValidProductCheckResponse().ResponseBytes
 					);
 				}
 				_ = Interlocked.Increment(ref state.Called);
@@ -139,7 +154,7 @@ public override TResponse Request<TResponse>(RequestData requestData)
 					nameof(VirtualCluster.ClientCalls),
 					_cluster.ClientCallRules,
 					requestData.RequestTimeout,
-					(r) => { },
+					_ => { },
 					CallResponse
 				);
 			}
@@ -247,14 +262,17 @@ private TResponse Fail<TResponse, TRule>(RequestData requestData, TRule rule, Ru
 				throw new TheException();
 
 			return ret.Match(
-				(e) => throw e,
-				(statusCode) => ReturnConnectionStatus<TResponse>(requestData, CallResponse(rule),
+				e => throw e,
+				statusCode => ReturnConnectionStatus<TResponse>(requestData, CallResponse(rule),
 					//make sure we never return a valid status code in Fail responses because of a bad rule.
-					statusCode >= 200 && statusCode < 300 ? 502 : statusCode, rule.ReturnContentType)
+					statusCode is >= 200 and < 300 ? 502 : statusCode, rule.ReturnContentType)
 			);
 		}
 
-		private TResponse Success<TResponse, TRule>(RequestData requestData, Action<TRule> beforeReturn, Func<TRule, byte[]> successResponse,
+		private TResponse Success<TResponse, TRule>(
+			RequestData requestData,
+			Action<TRule> beforeReturn,
+			Func<TRule, byte[]> successResponse,
 			TRule rule
 		)
 			where TResponse : class, IElasticsearchResponse, new()
@@ -277,11 +295,10 @@ private static byte[] CallResponse<TRule>(TRule rule)
 			if (_defaultResponseBytes != null) return _defaultResponseBytes;
 
 			var response = DefaultResponse;
-			using (var ms = RecyclableMemoryStreamFactory.Default.Create())
-			{
-				LowLevelRequestResponseSerializer.Instance.Serialize(response, ms);
-				_defaultResponseBytes = ms.ToArray();
-			}
+
+			using var ms = RecyclableMemoryStreamFactory.Default.Create();
+			LowLevelRequestResponseSerializer.Instance.Serialize(response, ms);
+			_defaultResponseBytes = ms.ToArray();
 			return _defaultResponseBytes;
 		}
 
@@ -295,6 +312,7 @@ private class State
 			public int Pinged;
 			public int Sniffed;
 			public int Successes;
+			public int ProductChecked;
 		}
 	}
 }
diff --git a/src/Elasticsearch.Net.VirtualizedCluster/VirtualClusterWith.cs b/src/Elasticsearch.Net.VirtualizedCluster/VirtualClusterWith.cs
