JSRuntimeInvocationHandlerBase{TResult}.cs

namespace Bunit;

/// <summary>
/// Represents an invocation handler for <see cref="JSRuntimeInvocation"/> instances.
/// </summary>
public abstract class JSRuntimeInvocationHandlerBase<TResult> : IDisposable
{
	private readonly InvocationMatcher invocationMatcher;
	private TaskCompletionSource<TResult> completionSource;
	private Timer? timeoutTimer;
	private JSRuntimeInvocation? currentInvocation;
	private bool disposed;

	/// <summary>
	/// Gets a value indicating whether this handler is set up to handle calls to <c>InvokeVoidAsync(string, object[])</c>.
	/// </summary>
	public virtual bool IsVoidResultHandler { get; }

	/// <summary>
	/// Gets a value indicating whether this handler is considered a catch all handler for invocations with <typeparamref name="TResult"/> as the return type.
	/// </summary>
	public bool IsCatchAllHandler { get; }

	/// <summary>
	/// Gets the invocations that this <see cref="JSRuntimeInvocationHandler{TResult}"/> has matched with.
	/// </summary>
	public JSRuntimeInvocationDictionary Invocations { get; } = new();

	/// <summary>
	/// Initializes a new instance of the <see cref="JSRuntimeInvocationHandlerBase{TResult}"/> class.
	/// </summary>
	/// <param name="matcher">An invocation matcher used to determine if the handler should handle an invocation.</param>
	/// <param name="isCatchAllHandler">Set to true if this handler is a catch all handler, that should only be used if there are no other non-catch all handlers available.</param>
	protected JSRuntimeInvocationHandlerBase(InvocationMatcher matcher, bool isCatchAllHandler)
	{
		invocationMatcher = matcher ?? throw new ArgumentNullException(nameof(matcher));
		completionSource = new TaskCompletionSource<TResult>(TaskCreationOptions.RunContinuationsAsynchronously);
		IsCatchAllHandler = isCatchAllHandler;
	}

	/// <summary>
	/// Marks the <see cref="Task{TResult}"/> that invocations will receive as canceled.
	/// </summary>
	protected void SetCanceledBase()
	{
		ClearTimeoutTimer();
		if (completionSource.Task.IsCompleted)
			completionSource = new TaskCompletionSource<TResult>(TaskCreationOptions.RunContinuationsAsynchronously);

		completionSource.SetCanceled();
	}

	/// <summary>
	/// Sets the <typeparamref name="TException"/> exception that invocations will receive.
	/// </summary>
	/// <param name="exception">The type of exception to pass to the callers.</param>
	protected void SetExceptionBase<TException>(TException exception)
		where TException : Exception
	{
		ClearTimeoutTimer();
		if (completionSource.Task.IsCompleted)
			completionSource = new TaskCompletionSource<TResult>(TaskCreationOptions.RunContinuationsAsynchronously);

		completionSource.SetException(exception);
	}

	/// <summary>
	/// Sets the <typeparamref name="TResult"/> result that invocations will receive.
	/// </summary>
	/// <param name="result">The type of result to pass to the callers.</param>
	protected void SetResultBase(TResult result)
	{
		ClearTimeoutTimer();
		if (completionSource.Task.IsCompleted)
			completionSource = new TaskCompletionSource<TResult>(TaskCreationOptions.RunContinuationsAsynchronously);

		completionSource.SetResult(result);
	}

	/// <summary>
	/// Call this to have the this handler handle the <paramref name="invocation"/>.
	/// </summary>
	/// <remarks>
	/// Note to implementors: Always call the <see cref="JSRuntimeInvocationHandlerBase{TResult}.HandleAsync(JSRuntimeInvocation)"/>
	/// method when overriding it in a sub class. It will make sure the invocation is correctly registered in the <see cref="Invocations"/> dictionary.
	/// </remarks>
	/// <param name="invocation">Invocation to handle.</param>
	protected internal virtual Task<TResult> HandleAsync(JSRuntimeInvocation invocation)
	{
		Invocations.RegisterInvocation(invocation);

		var task = completionSource.Task;
		if (task is { IsCanceled: false, IsFaulted: false, IsCompletedSuccessfully: false })
		{
			if (BunitContext.DefaultWaitTimeout <= TimeSpan.Zero)
			{
				throw new JSRuntimeInvocationNotSetException(invocation);
			}

			StartTimeoutTimer(invocation);
		}

		return task;
	}

	/// <summary>
	/// Checks whether this invocation handler can handle the <paramref name="invocation"/>.
	/// </summary>
	/// <param name="invocation">Invocation to check.</param>
	/// <returns>True if the handler can handle the invocation, false otherwise.</returns>
	internal bool CanHandle(JSRuntimeInvocation invocation) => invocationMatcher(invocation);

	/// <inheritdoc/>
	public void Dispose()
	{
		Dispose(disposing: true);
		GC.SuppressFinalize(this);
	}

	/// <inheritdoc/>
	protected virtual void Dispose(bool disposing)
	{
		if (!disposed && disposing)
		{
			ClearTimeoutTimer();
			disposed = true;
		}
	}

	private void StartTimeoutTimer(JSRuntimeInvocation invocation)
	{
		ClearTimeoutTimer();

		currentInvocation = invocation;
		timeoutTimer = new Timer(OnTimeoutElapsed, null, BunitContext.DefaultWaitTimeout, Timeout.InfiniteTimeSpan);
	}

	private void ClearTimeoutTimer()
	{
		timeoutTimer?.Dispose();
		timeoutTimer = null;
		currentInvocation = null;
	}

	private void OnTimeoutElapsed(object? state)
	{
		if (!completionSource.Task.IsCompleted && currentInvocation.HasValue)
		{
			var exception = new JSRuntimeInvocationNotSetException(currentInvocation.Value);
			completionSource.TrySetException(exception);
		}

		ClearTimeoutTimer();
	}
}