Support keyword arguments, mostly

This implements the last solution outlined in #7, and drops keywords
from the call when a recurse happens.
For some reason, it's faster than I remember this approach being.

Project.toml

@@ -1,7 +1,7 @@
 name = "SimpleMock"
 uuid = "a896ed2c-15a5-4479-b61d-a0e88e2a1d25"
 authors = ["Chris de Graaf <[email protected]>"]
-version = "0.3.2"
+version = "0.4.0"
 
 [deps]
 Cassette = "7057c7e9-c182-5462-911a-8362d720325c"

src/SimpleMock.jl

@@ -5,9 +5,9 @@ module SimpleMock
 
 using Base: Callable, invokelatest, unwrap_unionall
 using Base.Iterators: Pairs
-using Core: Builtin, IntrinsicFunction
+using Core: Builtin, IntrinsicFunction, kwftype
 
-using Cassette: Cassette, overdub, posthook, prehook, recurse, @context
+using Cassette: Cassette, Context, overdub, posthook, prehook, recurse, @context
 
 export
     Call,

src/mock_fun.jl

@@ -1,26 +1,11 @@
-@context MockCtx
-
-# TODO: Maybe these should be inlined, but it slows down compilation a lot.
-
-@noinline function Cassette.prehook(ctx::MockCtx{Metadata{true}}, f, args...)
-    @nospecialize f args
-    update!(ctx.metadata, prehook, f, args...)
-end
-
-@noinline function Cassette.posthook(ctx::MockCtx{Metadata{true}}, v, f, args...)
-    @nospecialize v f args
-    update!(ctx.metadata, posthook, f, args...)
-end
-
 """
-    mock(f::Function, args...; filters::Vector{<:Function}=Function[])
+    mock(f::Function[, ctx::Symbol], args...; filters::Vector{<:Function}=Function[])
 
 Run `f` with specified functions mocked out.
 
 !!! note
-    Keyword arguments to mocked functions are not supported.
-    If you call a mocked function with keyword arguments, it will dispatch to the original function.
-    For more details, see [Cassette#48](https://github.com/jrevels/Cassette.jl/issues/48).
+    Mocking functions with keyword arguments is only partially supported.
+    See the "Keyword Arguments" section below for more details.
 
 ## Examples
 
@@ -71,8 +56,8 @@ To avoid this, you can use filter functions like so:
 ```julia
 f(x, y) = x + y
 g(x, y) = f(x, y)
-mock((+) => Mock(; side_effect=(a, b) -> 2a + 2b); filters=[max_depth(2)]) do plus
-    @assert f(1, 2) == 6  # The call depth of + here is 2.
+mock((+) => Mock(; side_effect=(a, b) -> 2a + b); filters=[max_depth(2)]) do plus
+    @assert f(1, 2) == 4  # The call depth of + here is 2.
     @assert g(3, 4) == 7  # Here, it's 3.
     @assert called_once_with(plus, 1, 2)
 end
@@ -81,24 +66,75 @@ end
 Filter functions take a single argument of type [`Metadata`](@ref).
 If any filter rejects, then mocking is not performed.
 See [Filter Functions](@ref) for a list of included filters, as well as building blocks for you to create your own.
+
+## Reusing `Context`s
+
+Under the hood, this function creates a new [Cassette `Context`](https://jrevels.github.io/Cassette.jl/stable/api.html#Cassette.Context) on every call by default.
+This provides a nice clean mocking environment, but it can be slow to create and call new types and methods over and over.
+If you find yourself repeatedly mocking the same set of functions, you can specify a context name to reuse that context like so:
+
+```julia
+ctx = gensym()
+mock(g -> @assert(!called(g)), ctx, get)
+# This one is faster, especially when there's a lot going on in your mock blocks.
+mock(g -> @assert(!called(g)), ctx, get)
+```
+
+## Keyword Arguments
+
+Mocking of functions with keyword arguments is fully supported when the following conditions are met:
+
+- Filter functions are not used
+- The context in use has no previously-mocked functions that are now unmocked
+
+if a filter function rejects and calls a mocked function (instead of its mock), that call will have no keywords.
+
+If you reuse a context that has previously mocked some function, unmocked calls to that function will have no keywords.
+For example:
+
+```julia
+kwfunc(; kwargs...) = nothing
+calls_kwfunc(; kwargs...) = kwfunc(; kwargs...)
+ctx = gensym()
+mock(ctx, calls_kwfunc) do c
+    calls_kwfunc(; x=1, y=2)
+    @assert called_once_with(c; x=1, y=2)
+end
+mock(ctx, kwfunc) do k
+    calls_kwfunc(; x=1, y=2)               # This will issue a warning.
+    @assert called_once_with(k; x=1, y=2)  # This will fail!
+end
+```
+
+In short, avoid using filters and reusing contexts when mocking functions that accept keywords.
 """
-function mock(f::Function, args...; filters::Vector{<:Function}=Function[])
+mock(f::Function, args...; filters::Vector{<:Function}=Function[]) =
+    mock(f, gensym(), args...; filters=filters)
+
+function mock(f::Function, ctx::Symbol, args...; filters::Vector{<:Function}=Function[])
     mocks = map(sig2mock, args)  # ((f, sig) => mock).
     isempty(mocks) && throw(ArgumentError("At least one function must be mocked"))
 
+    # Create the new context type if it doesn't already exist.
+    context_is_new = !isdefined(@__MODULE__, ctx)
+    context_is_new && make_context(ctx)
+    Ctx = getfield(@__MODULE__, ctx)
+
     # Implement the overdubs, but only if they aren't already implemented.
     has_new_overdub = false
     foreach(map(first, mocks)) do k
         fun = k[1]
         sig = k[2:end]
-        if !overdub_exists(fun, sig)
-            make_overdub(fun, sig)
+        if context_is_new || !overdub_exists(Ctx, fun, sig)
+            make_overdub(Ctx, fun, sig)
             has_new_overdub = true
         end
     end
 
     # Only use `invokelatest` if the Context/overdub implementations are new.
-    od_args = [MockCtx(; metadata=Metadata(Dict(mocks), filters)), f, map(last, mocks)...]
+    meta = Metadata(Dict(mocks), filters)
+    c = context_is_new ? invokelatest(Ctx; metadata=meta) : Ctx(; metadata=meta)
+    od_args = [c, f, map(last, mocks)...]
     return has_new_overdub ? invokelatest(overdub, od_args...) : overdub(od_args...)
 end
 
@@ -108,20 +144,37 @@ sig2mock(p::Pair) = (p.first, Vararg{Any}) => p.second
 sig2mock(t::Tuple) = t => Mock()
 sig2mock(f) = (f, Vararg{Any}) => Mock()
 
+# Create a new context type.
+make_context(Ctx::Symbol) = @eval begin
+    @context $Ctx
+
+    # TODO: Maybe these should be inlined, but it slows down compilation a lot.
+
+    @noinline function Cassette.prehook(ctx::$Ctx{Metadata{true}}, f, args...)
+        @nospecialize f args
+        update!(ctx.metadata, prehook, f, args...)
+    end
+
+    @noinline function Cassette.posthook(ctx::$Ctx{Metadata{true}}, v, f, args...)
+        @nospecialize v f args
+        update!(ctx.metadata, posthook, f, args...)
+    end
+end
+
 # Has a given function and signature already been overdubbed?
-overdub_exists(::F, sig::Tuple) where F = any(methods(overdub)) do m
+function overdub_exists(::Type{Ctx}, ::F, sig::Tuple) where {Ctx <: Context, F}
     squashed = foldl(sig; init=[]) do acc, T
         if T isa DataType && T.name.name === :Vararg
             append!(acc, repeat([T.parameters[1]], T.parameters[2]))
         else
             push!(acc, T)
         end
     end
-    m.sig === Tuple{typeof(overdub), MockCtx, F, squashed...}
+    return any(m -> m.sig === Tuple{typeof(overdub), Ctx, F, squashed...}, methods(overdub))
 end
 
 # Implement `overdub` for a given Context, function, and signature.
-function make_overdub(f::F, sig::Tuple) where F
+function make_overdub(::Type{Ctx}, f::F, sig::Tuple) where {Ctx <: Context, F}
     sig_exs = Expr[]
     sig_names = []
 
@@ -149,12 +202,19 @@ function make_overdub(f::F, sig::Tuple) where F
         end
     end
 
-    @eval @inline function Cassette.overdub(ctx::MockCtx, f::$F, $(sig_exs...))
-        method = (f, $(sig...))
-        if should_mock(ctx.metadata, method)
-            ctx.metadata.mocks[method]($(sig_names...))
-        else
-            recurse(ctx, f, $(sig_names...))
+    @eval begin
+        @inline function Cassette.overdub(ctx::$Ctx, f::$F, $(sig_exs...); kwargs...)
+            method = (f, $(sig...))
+            if should_mock(ctx.metadata, method)
+                ctx.metadata.mocks[method]($(sig_names...); kwargs...)
+            else
+                isempty(kwargs) || @warn "Discarding keyword arguments" f kwargs
+                recurse(ctx, f, $(sig_names...))
+            end
         end
+
+        # https://github.com/jrevels/Cassette.jl/issues/48#issuecomment-440605481
+        @inline Cassette.overdub(ctx::$Ctx, ::kwftype($F), kwargs, f::$F, $(sig_exs...)) =
+            overdub(ctx, f, $(sig_names...); kwargs...)
     end
 end

test/mock_fun.jl

@@ -1,29 +1,12 @@
-@testset "mock does not overwrite methods" begin
-    # https://github.com/fredrikekre/jlpkg/blob/3b1c2400932dbe13fa7c3cba92bde3842315976c/src/cli.jl#L151-L160
-    o = JLOptions()
-    if o.warn_overwrite == 0
-        args = map(n -> n === :warn_overwrite ? 1 : getfield(o, n), fieldnames(JLOptions))
-        unsafe_store!(cglobal(:jl_options, JLOptions), JLOptions(args...))
-    end
-    mock(identity, identity)
-    out = @capture_err mock(identity, identity)
-    @test isempty(out)
-end
-
-@testset "Reusing Context" begin
-    f(x) = strip(uppercase(x))
-    # If the method checks aren't working properly, this will throw.
-    @test mock(_g -> f(" hi "), strip => identity) == " HI "
-    @test mock(_g -> f(" hi "), uppercase => identity) == "hi"
-end
+const IDENTITY_VA = gensym()
 
 @testset "Basics" begin
-    mock(identity) do id
+    mock(IDENTITY_VA, identity) do id
         identity(10)
         @test called_once_with(id, 10)
     end
 
-    mock(identity) do id
+    mock(IDENTITY_VA, identity) do id
         identity(1, 2, 3)
         identity()
         @test called(id)
@@ -58,12 +41,12 @@ end
 
 @testset "Varargs" begin
     varargs(::Int, ::Int, ::String, ::String, ::String, ::Bool...) = true
-    varargs(args...) = false
+    varargs(::Any) = false
 
     mock((varargs, Vararg{Int, 2}, Vararg{String, 3}, Vararg{Bool})) do va
         @test varargs(0, 0, "", "", "") !== true
         @test varargs(0, 0, "", "", "", false, false) !== true
-        @test !varargs()
+        @test !varargs(0)
         @test ncalls(va) == 2
         @test has_calls(va, Call(0, 0, "", "", ""), Call(0, 0, "", "", "", false, false))
     end
@@ -96,20 +79,41 @@ end
     end
 end
 
+@testset "mock does not overwrite methods" begin
+    # https://github.com/fredrikekre/jlpkg/blob/3b1c2400932dbe13fa7c3cba92bde3842315976c/src/cli.jl#L151-L160
+    o = JLOptions()
+    if o.warn_overwrite == 0
+        args = map(n -> n === :warn_overwrite ? 1 : getfield(o, n), fieldnames(JLOptions))
+        unsafe_store!(cglobal(:jl_options, JLOptions), JLOptions(args...))
+    end
+    ctx = gensym()
+    mock(identity, ctx, identity)
+    out = @capture_err mock(identity, ctx, identity)
+    @test isempty(out)
+end
+
+@testset "Reusing Context" begin
+    f(x) = strip(uppercase(x))
+    # If the method checks aren't working properly, this will throw.
+    ctx = gensym()
+    @test mock(_f -> f(" hi "), ctx, strip => identity) == " HI "
+    @test mock(_f -> f(" hi "), ctx, uppercase => identity) == "hi"
+end
+
 @testset "Filters" begin
     @testset "Maximum/minimum depth" begin
         f(x) = identity(x)
         g(x) = f(x)
         h(x) = g(x)
 
-        mock(identity; filters=[max_depth(3)]) do id
+        mock(IDENTITY_VA, identity; filters=[max_depth(3)]) do id
             @test f(1) != 1
             @test g(2) != 2
             @test h(3) == 3
             @test ncalls(id) == 2 && has_calls(id, Call(1), Call(2))
         end
 
-        mock(identity; filters=[min_depth(3)]) do id
+        mock(IDENTITY_VA, identity; filters=[min_depth(3)]) do id
             @test f(1) == 1
             @test g(2) != 2
             @test h(3) != 3
@@ -125,18 +129,47 @@ end
         c(x) = identity(x)
         d(x) = identity(x)
 
-        mock(identity; filters=[excluding(Bar, c)]) do id
+        mock(IDENTITY_VA, identity; filters=[excluding(Bar, c)]) do id
             @test Bar.a(1) == 1
             @test Bar.b(2) == 2
             @test c(3) == 3
             @test d(4) != 4
         end
 
-        mock(identity; filters=[including(Bar, c)]) do id
+        mock(IDENTITY_VA, identity; filters=[including(Bar, c)]) do id
             @test Bar.a(1) != 1
             @test Bar.b(2) != 2
             @test c(3) != 3
             @test d(4) == 4
         end
     end
 end
+
+@testset "Keyword arguments" begin
+    foo(; kwargs...) = get(kwargs, :foo, nothing)
+    bar(; kwargs...) = foo(; kwargs...)
+    baz(; kwargs...) = bar(; kwargs...)
+
+    @testset "Keyword arguments are passed to mocked functions" begin
+        mock(foo) do f
+            @test bar(; foo=:bar) !== :bar
+            @test called_once_with(f; foo=:bar)
+        end
+    end
+
+    @testset "Keyword arguments are discarded when recursing" begin
+        ctx = gensym()
+
+        mock(ctx, bar; filters=[including()]) do _b
+            @test_logs (:warn, "Discarding keyword arguments") baz(; foo=:baz)
+            @test @suppress baz(; foo=:baz) === nothing
+        end
+
+        mock(ctx, foo) do f
+            @test_logs (:warn, "Discarding keyword arguments") baz(; foo=:baz)
+            result = @suppress baz(; foo=:baz)
+            @test result !== nothing && result !== :baz
+            @test ncalls(f) == 2 && has_calls(f, Call(), Call())
+        end
+    end
+end

test/runtests.jl

@@ -1,8 +1,8 @@
 using Base: JLOptions
 
-using Test: @test, @testset, @test_throws
+using Test: @test, @testset, @test_logs, @test_throws
 
-using Suppressor: @capture_err
+using Suppressor: @capture_err, @suppress
 
 using SimpleMock