MemoryProfilerRule.kt

class MemoryProfilerRule : TestRule {
    override fun apply(base: Statement?, description: Description?): Statement {
        return MemoryProfilerWatcher().apply(base, description)
    }

    private inner class MemoryProfilerWatcher : TestWatcher() {

        private val runtime = Runtime.getRuntime()

        private var before: Long = 0
        private var after: Long = 0

        override fun starting(description: Description?) {
            before = getCurrentlyAllocatedMemory()
        }

        override fun finished(description: Description?) {
            after = getCurrentlyAllocatedMemory()

            println("""
                Before: $before
                After: $after
            """.trimIndent())
        }

        fun getCurrentlyAllocatedMemory(): Long {
            return (runtime.totalMemory() - runtime.freeMemory()) / (1024 * 1024)
        }
    }
}