ShouldlyJsonExtensions.cs

using System.Text.Json;
using System.Text.Json.Nodes;
using PropertyDiff = (string Path, object? ExpectedValue, object? ActualValue);

namespace Tests;

[ShouldlyMethods]
public static class ShouldlyJsonExtensions
{
    public static void ShouldBeEquivalentTo(this JsonNode actual, JsonNode expected, string? customMessage = null)
    {
        if (JsonNode.DeepEquals(expected, actual))
        {
            return;
        }

        var diff = GetFirstDiff(expected, actual);
        if (diff is null)
        {
            throw new ShouldAssertException(new ExpectedEquivalenceShouldlyMessage(expected.ToJsonString(), actual.ToJsonString(), ["<root>"], customMessage).ToString());
        }

        throw new ShouldAssertException(
            new ExpectedEquivalenceShouldlyMessage(diff.Value.ExpectedValue, diff.Value.ActualValue, diff.Value.Path.Split("."), customMessage).ToString());
    }

    private static PropertyDiff? GetFirstDiff(JsonNode expected, JsonNode actual)
    {
        if (actual.GetValueKind() != expected.GetValueKind())
        {
            return new PropertyDiff("<root>", expected, actual);
        }

        if (actual.GetValueKind() is not (JsonValueKind.Object or JsonValueKind.Array))
        {
            return Equals(actual.GetValue(), expected.GetValue()) ? null : new PropertyDiff("<root>", expected, actual);
        }

        var expectedFlattened = expected.Flatten();
        var actualFlattened = actual.Flatten();

        var expectedKeys = expectedFlattened.Keys.ToHashSet();
        var actualKeys = actualFlattened.Keys.ToHashSet();

        // get paths that are either in expected or in both
        var paths = new HashSet<string>(expectedKeys);
        paths.IntersectWith(actualKeys);
        paths.UnionWith(expectedKeys);

        foreach (var path in paths)
        {
            if (!Equals(actualFlattened.GetValueOrDefault(path), expectedFlattened.GetValueOrDefault(path)))
            {
                return new PropertyDiff(path, expectedFlattened.GetValueOrDefault(path), actualFlattened.GetValueOrDefault(path));
            }
        }

        return null;
    }

    private static Dictionary<string, object?> Flatten(this JsonNode json)
    {
        return json.GetValueKind() switch
        {
            JsonValueKind.Object => json.AsObject()
                .SelectMany(p => GetLeaves(null!, $".{p.Key}", p.Value))
                .ToDictionary(k => k.Path, v => v.Element.AsValue().GetValue()),

            JsonValueKind.Array => json.AsArray()
                .SelectMany((p, idx) => GetLeaves(null!, $".[{idx}]", p))
                .ToDictionary(k => k.Path, v => v.Element.AsValue().GetValue()),

            _ => new Dictionary<string, object?> { ["<root>"] = json.GetValue() }
        };
    }

    private static object? GetValue(this JsonNode element)
    {
        return element.GetValueKind() switch
        {
            JsonValueKind.True => true,
            JsonValueKind.False => false,
            JsonValueKind.String => element.GetValue<string>(),
            JsonValueKind.Number => GetNumericValue(element),
            JsonValueKind.Null or JsonValueKind.Undefined => null,
            JsonValueKind.Array => element.Deserialize<IEnumerable<object>>(),
            JsonValueKind.Object => element.Deserialize<Dictionary<string, object>>(),
            _ => throw new ArgumentOutOfRangeException("Unknown value kind: " + element.GetValueKind())
        };
    }

    private static object GetNumericValue(JsonNode element)
    {
        try
        {
            return element.GetValue<int>();
        }
        catch
        {
            return element.GetValue<double>();
        }
    }

    private static IEnumerable<(string Path, JsonNode Element)> GetLeaves(string? path, string name, JsonNode? element)
    {
        if (element is null)
        {
            return [];
        }

        return element.GetValueKind() switch
        {
            JsonValueKind.Object => element
                .AsObject()
                .SelectMany(child => GetLeaves(path == null ? $"{name}" : $"{path}.{name}", child.Key, child.Value)),

            JsonValueKind.Array => element.AsArray()
                .SelectMany((child, idx) =>
                    child!.GetValueKind() is JsonValueKind.Object
                        ? child.AsObject()
                            .SelectMany(innerChild => GetLeaves(GetArrayPath(path, name, idx), innerChild.Key, innerChild.Value))
                        : new[] { (Path: GetArrayPath(path, name, idx), child) }
                ),

            _ => [(Path: path == null ? $"{name}" : $"{path}.{name}", element)]
        };
    }

    private static string GetArrayPath(string? path, string name, int idx) => path == null ? $"{name}.[{idx}]" : $"{path}.{name}.[{idx}]";
}