feat: Add FindComponent on IElement (#153)

src/bunit.generators.internal/Web.AngleSharp/IElementWrapper.cs

@@ -14,5 +14,10 @@ public interface IElementWrapper<out TElement> where TElement : class, IElement
 	/// Gets the wrapped element.
 	/// </summary>
 	TElement WrappedElement { get; }
+
+	/// <summary>
+	/// Gets the element wrapper factory used by this wrapper.
+	/// </summary>
+	IElementWrapperFactory Factory { get; }
 }
 #nullable restore

src/bunit.generators.internal/Web.AngleSharp/WrapperBase.cs

@@ -33,6 +33,16 @@ public TElement WrappedElement
 		}
 	}
 
+	/// <summary>
+	/// Gets the element wrapper factory used by this wrapper.
+	/// </summary>
+	[DebuggerNonUserCode]
+	public IElementWrapperFactory Factory
+	{
+		[MethodImpl(MethodImplOptions.AggressiveInlining)]
+		get => elementFactory;
+	}
+
 	/// <summary>
 	/// Creates an instance of the <see cref="WrapperBase{T}"/> class.
 	/// </summary>

src/bunit.web.query/ByLabelTextElementFactory.cs

@@ -3,14 +3,16 @@
 
 namespace Bunit;
 
-internal sealed class ByLabelTextElementFactory : IElementWrapperFactory
+internal sealed class ByLabelTextElementFactory : IElementWrapperFactory, IComponentAccessor
 {
 	private readonly IRenderedComponent<IComponent> testTarget;
 	private readonly string labelText;
 	private readonly ByLabelTextOptions options;
 
 	public Action? OnElementReplaced { get; set; }
 
+	public IRenderedComponent<IComponent> Component => testTarget;
+
 	public ByLabelTextElementFactory(IRenderedComponent<IComponent> testTarget, string labelText, ByLabelTextOptions options)
 	{
 		this.testTarget = testTarget;

src/bunit/Extensions/ElementExtensions.cs

@@ -0,0 +1,44 @@
+using AngleSharp.Dom;
+using Bunit.Rendering;
+
+namespace Bunit;
+
+/// <summary>
+/// Provides extension methods for <see cref="IElement"/> to find components rendered inside them.
+/// </summary>
+public static class ElementExtensions
+{
+	/// <summary>
+	/// Retrieves the first component of type <typeparamref name="TComponent"/> that is rendered inside the specified <paramref name="element"/>.
+	/// </summary>
+	public static IRenderedComponent<TComponent> FindComponent<TComponent>(this IElement element)
+		where TComponent : IComponent
+	{
+		ArgumentNullException.ThrowIfNull(element);
+
+		var componentAccessor = element.GetComponentAccessor();
+		if (componentAccessor is null)
+		{
+			throw new InvalidOperationException(
+				$"Unable to find component of type {typeof(TComponent).Name} for the given element.");
+		}
+
+		var component = componentAccessor.Component;
+		if (component.Instance is TComponent)
+		{
+			return (IRenderedComponent<TComponent>)component;
+		}
+
+		var renderer = GetRendererFromComponent(component);
+		var foundComponent = renderer.FindComponentForElement<TComponent>(element);
+		if (foundComponent is not null)
+		{
+			return foundComponent;
+		}
+
+		throw new InvalidOperationException($"Unable to find component of type {typeof(TComponent).Name} for the given element.");
+	}
+
+	private static BunitRenderer GetRendererFromComponent(IRenderedComponent<IComponent> component)
+		=> component.Services.GetRequiredService<BunitContext>().Renderer;
+}

src/bunit/Extensions/Internal/CssSelectorElementFactory.cs

@@ -3,13 +3,15 @@
 
 namespace Bunit;
 
-internal sealed class CssSelectorElementFactory : IElementWrapperFactory
+internal sealed class CssSelectorElementFactory : IElementWrapperFactory, IComponentAccessor
 {
 	private readonly IRenderedComponent<IComponent> testTarget;
 	private readonly string cssSelector;
 
 	public Action? OnElementReplaced { get; set; }
 
+	public IRenderedComponent<IComponent> Component => testTarget;
+
 	public CssSelectorElementFactory(IRenderedComponent<IComponent> testTarget, string cssSelector)
 	{
 		this.testTarget = testTarget;

src/bunit/Extensions/Internal/ElementFactoryAccessor.cs

@@ -0,0 +1,26 @@
+using AngleSharp.Dom;
+using Bunit.Web.AngleSharp;
+
+namespace Bunit;
+
+/// <summary>
+/// Internal extensions for working with wrapped elements.
+/// </summary>
+internal static class ElementFactoryAccessor
+{
+	/// <summary>
+	/// Attempts to get the component accessor from an element factory.
+	/// </summary>
+	/// <param name="element">The element to get the component from.</param>
+	/// <returns>The component accessor if available, null otherwise.</returns>
+	internal static IComponentAccessor? GetComponentAccessor(this IElement element)
+	{
+		var factory = element.GetElementFactory();
+		return factory as IComponentAccessor;
+	}
+
+	private static IElementWrapperFactory? GetElementFactory(this IElement element)
+	{
+		return element is IElementWrapper<IElement> wrapper ? wrapper.Factory : null;
+	}
+}

src/bunit/Extensions/Internal/IComponentAccessor.cs

@@ -0,0 +1,12 @@
+namespace Bunit;
+
+/// <summary>
+/// Interface for accessing the component that owns an element.
+/// </summary>
+internal interface IComponentAccessor
+{
+	/// <summary>
+	/// Gets the component that owns this element.
+	/// </summary>
+	IRenderedComponent<IComponent> Component { get; }
+}

src/bunit/InternalsVisibleTo.cs

@@ -1 +1,2 @@
 [assembly: System.Runtime.CompilerServices.InternalsVisibleTo("Bunit.Tests, PublicKey=002400000480000094000000060200000024000052534131000400000100010001be6b1a2ca57b09b7040e2ab0993e515296ae22aef4031a4fe388a1336fe21f69c7e8610e9935de6ed18d94b5c98429f99ef62ce3d0af28a7088f856239368ea808ad4c448aa2a8075ed581f989f36ed0d0b8b1cfcaf1ff6a4506c8a99b7024b6eb56996d08e3c9c1cf5db59bff96fcc63ccad155ef7fc63aab6a69862437b6")]
+[assembly: System.Runtime.CompilerServices.InternalsVisibleTo("Bunit.Web.Query, PublicKey=002400000480000094000000060200000024000052534131000400000100010001be6b1a2ca57b09b7040e2ab0993e515296ae22aef4031a4fe388a1336fe21f69c7e8610e9935de6ed18d94b5c98429f99ef62ce3d0af28a7088f856239368ea808ad4c448aa2a8075ed581f989f36ed0d0b8b1cfcaf1ff6a4506c8a99b7024b6eb56996d08e3c9c1cf5db59bff96fcc63ccad155ef7fc63aab6a69862437b6")]

src/bunit/Rendering/BunitRenderer.cs

@@ -3,6 +3,7 @@
 using System.Reflection;
 using System.Runtime.CompilerServices;
 using System.Runtime.ExceptionServices;
+using AngleSharp.Dom;
 using Microsoft.Extensions.Logging;
 
 namespace Bunit.Rendering;
@@ -26,6 +27,7 @@ public sealed class BunitRenderer : Renderer
 
 	private readonly HashSet<int> returnedRenderedComponentIds = new();
 	private readonly List<BunitRootComponent> rootComponents = new();
+	private readonly Dictionary<string, int> elementReferenceToComponentId = new();
 	private readonly ILogger<BunitRenderer> logger;
 	private bool disposed;
 	private TaskCompletionSource<Exception> unhandledExceptionTsc = new(TaskCreationOptions.RunContinuationsAsynchronously);
@@ -453,6 +455,7 @@ protected override Task UpdateDisplayAsync(in RenderBatch renderBatch)
 			var id = renderBatch.DisposedComponentIDs.Array[i];
 			disposedComponentIds.Add(id);
 			returnedRenderedComponentIds.Remove(id);
+			RemoveElementReferencesForComponent(id);
 		}
 
 		for (var i = 0; i < renderBatch.UpdatedComponents.Count; i++)
@@ -467,6 +470,8 @@ protected override Task UpdateDisplayAsync(in RenderBatch renderBatch)
 			var componentState = GetComponentState(diff.ComponentId);
 			var renderedComponent = (IRenderedComponent)componentState;
 
+			TrackElementReferencesForComponent(diff.ComponentId);
+
 			if (returnedRenderedComponentIds.Contains(diff.ComponentId))
 			{
 				renderedComponent.UpdateState(hasRendered: true, isMarkupGenerationRequired: diff.Edits.Count > 0);
@@ -519,6 +524,101 @@ static bool IsParentComponentAlreadyUpdated(int componentId, in RenderBatch rend
 		}
 	}
 
+	private void TrackElementReferencesForComponent(int componentId)
+	{
+		var frames = GetCurrentRenderTreeFrames(componentId);
+		TrackElementReferencesInFrames(frames, componentId);
+	}
+
+	private void TrackElementReferencesInFrames(ArrayRange<RenderTreeFrame> frames, int componentId)
+	{
+		for (var i = 0; i < frames.Count; i++)
+		{
+			ref var frame = ref frames.Array[i];
+
+			if (frame.FrameType == RenderTreeFrameType.ElementReferenceCapture)
+			{
+				var elementReferenceId = frame.ElementReferenceCaptureId;
+				if (elementReferenceId != null)
+				{
+					elementReferenceToComponentId[elementReferenceId] = componentId;
+				}
+			}
+			else if (frame.FrameType == RenderTreeFrameType.Component)
+			{
+				TrackElementReferencesForComponent(frame.ComponentId);
+			}
+		}
+	}
+
+	private void RemoveElementReferencesForComponent(int componentId)
+	{
+		var keysToRemove = elementReferenceToComponentId
+			.Where(kvp => kvp.Value == componentId)
+			.Select(kvp => kvp.Key)
+			.ToList();
+
+		foreach (var key in keysToRemove)
+		{
+			elementReferenceToComponentId.Remove(key);
+		}
+	}
+
+	internal IRenderedComponent<TComponent>? FindComponentForElement<TComponent>(IElement element)
+		where TComponent : IComponent
+	{
+		var elementReferenceId = element.GetAttribute("blazor:elementReference");
+		if (elementReferenceId is not null && elementReferenceToComponentId.TryGetValue(elementReferenceId, out var componentId))
+		{
+			return GetRenderedComponent<TComponent>(componentId);
+		}
+
+		return FindComponentByElementContainment<TComponent>(element);
+	}
+
+	private IRenderedComponent<TComponent>? FindComponentByElementContainment<TComponent>(IElement element)
+		where TComponent : IComponent
+	{
+		List<int> renderedComponentIdsWhenStarted = [..returnedRenderedComponentIds];
+		var components = new List<IRenderedComponent<TComponent>>(returnedRenderedComponentIds.Count);
+
+		foreach (var parentComponent in renderedComponentIdsWhenStarted.Select(GetRenderedComponent<IComponent>))
+		{
+			components.AddRange(FindComponents<TComponent>(parentComponent));
+		}
+
+		return components.FirstOrDefault(component => ComponentContainsElement(component, element));
+	}
+
+	private static bool ComponentContainsElement<TComponent>(IRenderedComponent<TComponent> component, IElement element)
+		where TComponent : IComponent
+	{
+		foreach (var node in component.Nodes)
+		{
+			if (node is IElement nodeElement && nodeElement.Equals(element))
+			{
+				return true;
+			}
+			if (IsDescendantOf(element, node))
+			{
+				return true;
+			}
+		}
+		return false;
+	}
+
+	private static bool IsDescendantOf(IElement element, INode potentialAncestor)
+	{
+		var current = element.Parent;
+		while (current is not null)
+		{
+			if (current == potentialAncestor)
+				return true;
+			current = current.Parent;
+		}
+		return false;
+	}
+
 	/// <inheritdoc/>
 	internal new ArrayRange<RenderTreeFrame> GetCurrentRenderTreeFrames(int componentId)
 		=> base.GetCurrentRenderTreeFrames(componentId);