1 files changed, 37 insertions, 3 deletions
diff --git a/core/src/main/kotlin/Samples/KotlinWebsiteSampleProcessingService.kt b/core/src/main/kotlin/Samples/KotlinWebsiteSampleProcessingService.kt
index 7ac43184..13b9f2de 100644
--- a/core/src/main/kotlin/Samples/KotlinWebsiteSampleProcessingService.kt
+++ b/core/src/main/kotlin/Samples/KotlinWebsiteSampleProcessingService.kt
@@ -3,6 +3,7 @@ package org.jetbrains.dokka.Samples
 import com.google.inject.Inject
 import com.intellij.psi.PsiElement
 import com.intellij.psi.impl.source.tree.LeafPsiElement
+import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.dokka.*
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.allChildren
@@ -19,15 +20,18 @@ open class KotlinWebsiteSampleProcessingService
         val text: String
             get() = builder.toString()
 
+        fun KtValueArgument.extractStringArgumentValue() =
+                (getArgumentExpression() as KtStringTemplateExpression)
+                        .entries.joinToString("") { it.text }
+
+
         fun convertAssertPrints(expression: KtCallExpression) {
             val (argument, commentArgument) = expression.valueArguments
-            val comment = commentArgument.getArgumentExpression() as KtStringTemplateExpression
-            val commentText = comment.entries.joinToString("") { it.text }
             builder.apply {
                 append("println(")
                 append(argument.text)
                 append(") // ")
-                append(commentText)
+                append(commentArgument.extractStringArgumentValue())
             }
         }
 
@@ -42,10 +46,40 @@ open class KotlinWebsiteSampleProcessingService
             }
         }
 
+        fun convertAssertFails(expression: KtCallExpression) {
+            val (message, funcArgument) = expression.valueArguments
+            builder.apply {
+                val argument = if (funcArgument.getArgumentExpression() is KtLambdaExpression)
+                    PsiTreeUtil.findChildOfType(funcArgument, KtBlockExpression::class.java)?.text ?: ""
+                else
+                    funcArgument.text
+                append(argument.lines().joinToString(separator = "\n") { "// $it" })
+                append(" // ")
+                append(message.extractStringArgumentValue())
+                append(" will fail")
+            }
+        }
+
+        fun convertAssertFailsWith(expression: KtCallExpression) {
+            val (funcArgument) = expression.valueArguments
+            val (exceptionType) = expression.typeArguments
+            builder.apply {
+                val argument = if (funcArgument.firstChild is KtLambdaExpression)
+                    PsiTreeUtil.findChildOfType(funcArgument, KtBlockExpression::class.java)?.text ?: ""
+                else
+                    funcArgument.text
+                append(argument.lines().joinToString(separator = "\n") { "// $it" })
+                append(" // will fail with ")
+                append(exceptionType.text)
+            }
+        }
+
         override fun visitCallExpression(expression: KtCallExpression) {
             when (expression.calleeExpression?.text) {
                 "assertPrints" -> convertAssertPrints(expression)
                 "assertTrue" -> convertAssertTrue(expression)
+                "assertFails" -> convertAssertFails(expression)
+                "assertFailsWith" -> convertAssertFailsWith(expression)
                 else -> super.visitCallExpression(expression)
             }
         }