How to unit test function that calls Amplify Api mutate method?

I have a system that is using Amplify Api to interact with a GraphQL api that is not part of an Amplify project. None of it was set up with the amplify-cli. I have a function that makes a call to the mutate function of the AmplifyCategory class. I would like to unit test my function, making sure to test both callbacks of the mutate method. How can I mock the mutate method to tell it which callback to use, and test that the callback was, in fact, called?

MyClass.kt

import com.amplifyframework.api.ApiCategory
import com.amplifyframework.api.aws.GsonVariablesSerializer
import com.amplifyframework.api.graphql.SimpleGraphQLRequest

class MyClass {
    private val graphQlClient: ApiCategory = [... do stuff to configure client]

    suspend fun foo(): Boolean {
        val mutationDefinition = """
            mutation doFoo {
              doFoo
            }""";
        try {
            val mutation = SimpleGraphQLRequest<Boolean>(mutationDefinition, Boolean::class.java, GsonVariablesSerializer())
            var success: Boolean = false
            graphQlClient.mutate(mutation,
                {
                    success = it.data
                },
                {
                    logger.error("ERROR!")
                }
            )
            return success
        } catch (e: Throwable) {
            logger.debug("unexpected error $e")
            throw e
        }
    }
}

TestMyClass.kt

import com.amplifyframework.api.ApiCategory
import com.amplifyframework.api.graphql.GraphQLOperation
import com.amplifyframework.api.graphql.GraphQLRequest
import com.amplifyframework.api.graphql.GraphQLResponse
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.async
import kotlinx.coroutines.delay
import kotlinx.coroutines.runBlocking
import org.junit.Test
import org.mockito.kotlin.any
import org.mockito.kotlin.doReturn
import org.mockito.kotlin.mock
import org.mockito.kotlin.stub
import org.mockito.kotlin.verify

class TestMyClass {
    [A bunch of stuff to initialize...]
    private val mutationResult by before {
        GraphQLResponse(true, null)
    }

    private val graphQlOperation by before {
        mock<GraphQLOperation<Boolean>>()
    }

    private val mockGraphQlClient by before {
        mock<ApiCategory>().stub {
            on { mutate(any<GraphQLRequest<Boolean>>(), any(), any()) } doReturn graphQlOperation
        }
    }

    private val client by before {
        MyClass()
    }

    @Test
    fun `foo() should return true when mutate succeeds`() = runBlocking<Unit> {
        val deferredResult = async(Dispatchers.IO) {
            client.foo()
        }
        deferredResult.start()

        delay(500L)

        val result = deferredResult.await()
        result shouldBe true // Fails here, presumably because the success lambda never gets invoked

        verify(mockGraphQlClient).mutate(
            any<GraphQLRequest<Boolean>>(),
            any(),
            any()
        )
    }
}

I’ve tried a few different things that have not worked, including using Mockito’s whenever on the mutate function to try to invoke the second argument, but that fails because the lambda cannot be case to a GraphQlOperation

whenever(mockGraphQlClient.mutate(any<GraphQLRequest<Boolean>>(), any(), any())).thenAnswer {
    it.arguments[0].success()
}

Caused by: java.lang.ClassCastException: class arrow.core.Try$Success cannot be cast to class com.amplifyframework.api.graphql.GraphQLOperation (arrow.core.Try$Success and com.amplifyframework.api.graphql.GraphQLOperation are in unnamed module of loader 'app')

I’ve also tried using an argumentCaptor to capture the arguments being passed to mutate, but I can’t verify that they’ve been called because they aren’t mocks.

With all of these cases, I’m pretty sure the lambdas are not actually being invoked. Any help would be greatly appreciated.