Fix issues with block media at the end of editable text

aztec/src/main/kotlin/org/wordpress/aztec/formatting/LineBlockFormatter.kt

@@ -92,6 +92,7 @@ class LineBlockFormatter(editor: AztecText) : AztecFormatter(editor) {
 
         return false
     }
+
     fun applyHorizontalRule(inline: Boolean) {
         val nestingLevel = if (inline) {
             editor.removeInlineStylesFromRange(selectionStart, selectionEnd)
@@ -117,7 +118,6 @@ class LineBlockFormatter(editor: AztecText) : AztecFormatter(editor) {
             val newSelectionPosition = editableText.indexOf(Constants.MAGIC_CHAR, selectionStart) + 1
             editor.setSelection(newSelectionPosition)
         } else {
-            builder.append("\n")
             insertSpanAfterBlock(builder)
         }
     }
@@ -163,7 +163,6 @@ class LineBlockFormatter(editor: AztecText) : AztecFormatter(editor) {
 
     private fun insertMediaAfterBlock(span: AztecMediaSpan) {
         val ssb = SpannableStringBuilder(Constants.IMG_STRING)
-        ssb.append("\n")
         buildClickableMediaSpan(ssb, span)
         insertSpanAfterBlock(ssb)
     }
@@ -173,23 +172,39 @@ class LineBlockFormatter(editor: AztecText) : AztecFormatter(editor) {
         // We need to be sure the cursor is placed correctly after media insertion
         // Note that media has '\n' around them when needed
         val isLastItem = position == EndOfBufferMarkerAdder.safeLength(editor)
-        val insertedLength = ssb.length
-        editableText.insert(position, ssb)
-        val spans = editableText.getSpans(position, position + insertedLength, IAztecBlockSpan::class.java).filter {
-            it !is AztecMediaSpan && editableText.getSpanStart(it) == position
-        }
-        spans.forEach {
-            val spanStart = editableText.getSpanStart(it)
-            val spanEnd = editableText.getSpanEnd(it)
-            val spanFlags = editableText.getSpanFlags(it)
-            editableText.removeSpan(it)
-            if (spanStart + insertedLength < spanEnd) {
-                editableText.setSpan(it, spanStart + insertedLength, spanEnd, spanFlags)
+        if (isLastItem) {
+            editableText.getSpans(position, editableText.length, IAztecBlockSpan::class.java).filter {
+                it !is AztecMediaSpan && editableText.getSpanEnd(it) == editableText.length
+            }.map {
+                SpanData(it, editableText.getSpanStart(it), position + 1, editableText.getSpanFlags(it))
+            }.applyWithRemovedSpans {
+                editableText.append(ssb)
+            }
+        } else {
+            ssb.append("\n")
+
+            val ssbLength = ssb.length
+            editableText.getSpans(position, position + ssbLength, IAztecBlockSpan::class.java).filter {
+                it !is AztecMediaSpan && editableText.getSpanStart(it) == position
+            }.map {
+                SpanData(it, editableText.getSpanStart(it) + ssbLength, editableText.getSpanEnd(it) + ssbLength, editableText.getSpanFlags(it))
+            }.applyWithRemovedSpans {
+                editableText.insert(position, ssb)
             }
         }
         setSelection(isLastItem, position)
     }
 
+    private fun List<SpanData>.applyWithRemovedSpans(action: () -> Unit) {
+        this.onEach { editableText.removeSpan(it.span) }
+        action()
+        this.onEach {
+            editableText.setSpan(it.span, it.spanStart, it.spanEnd, it.spanFlags)
+        }
+    }
+
+    data class SpanData(val span: IAztecBlockSpan, val spanStart: Int, val spanEnd: Int, val spanFlags: Int)
+
     private fun setSelection(isLastItem: Boolean, position: Int) {
         val newSelection = if (isLastItem) {
             EndOfBufferMarkerAdder.safeLength(editor)

aztec/src/test/kotlin/org/wordpress/aztec/ImageBlockTest.kt

@@ -0,0 +1,140 @@
+package org.wordpress.aztec
+
+import android.app.Activity
+import android.view.MenuItem
+import android.widget.PopupMenu
+import org.junit.Assert
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.Robolectric
+import org.robolectric.RobolectricTestRunner
+import org.robolectric.annotation.Config
+import org.wordpress.aztec.source.SourceViewEditText
+import org.wordpress.aztec.toolbar.AztecToolbar
+
+@RunWith(RobolectricTestRunner::class)
+@Config(sdk = intArrayOf(23))
+class ImageBlockTest {
+    lateinit var editText: AztecText
+    lateinit var menuList: PopupMenu
+    lateinit var menuListOrdered: MenuItem
+    lateinit var menuListUnordered: MenuItem
+    lateinit var sourceText: SourceViewEditText
+    lateinit var toolbar: AztecToolbar
+
+    /**
+     * Initialize variables.
+     */
+    @Before
+    fun init() {
+        val activity = Robolectric.buildActivity(Activity::class.java).create().visible().get()
+        editText = AztecText(activity)
+        editText.setCalypsoMode(false)
+        editText.addMediaAfterBlocks()
+        sourceText = SourceViewEditText(activity)
+        sourceText.setCalypsoMode(false)
+        toolbar = AztecToolbar(activity)
+        toolbar.setEditor(editText, sourceText)
+        menuList = toolbar.getListMenu() as PopupMenu
+        menuListOrdered = menuList.menu.getItem(1)
+        menuListUnordered = menuList.menu.getItem(0)
+        activity.setContentView(editText)
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun addImageAfterAListAtTheEnd() {
+        editText.fromHtml("<ul><li>item 1</li><li>item 2</li></ul>")
+
+        editText.setSelection(editText.editableText.indexOf("2"))
+        val attributes = AztecAttributes()
+        attributes.setValue("id", "1234")
+        editText.insertImage(null, attributes)
+
+        Assert.assertEquals("<ul><li>item 1</li><li>item 2</li></ul><img id=\"1234\" />", editText.toHtml())
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun addHRAfterAListAtTheEnd() {
+        editText.fromHtml("<ul><li>item 1</li><li>item 2</li></ul>")
+
+        editText.setSelection(editText.editableText.indexOf("2"))
+        editText.lineBlockFormatter.applyHorizontalRule(false)
+
+        Assert.assertEquals("<ul><li>item 1</li><li>item 2</li></ul><hr />", editText.toHtml())
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun addImageAfterAListInTheMiddle() {
+        editText.fromHtml("<ul><li>item 1</li><li>item 2</li></ul>\n<p>test</p>")
+
+        editText.setSelection(editText.editableText.indexOf("2"))
+        val attributes = AztecAttributes()
+        attributes.setValue("id", "1234")
+        editText.insertImage(null, attributes)
+
+        Assert.assertEquals("<ul><li>item 1</li><li>item 2</li></ul><img id=\"1234\" /><p>test</p>", editText.toHtml())
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun addHRAfterAListInTheMiddle() {
+        editText.fromHtml("<ul><li>item 1</li><li>item 2</li></ul>\n<p>test</p>")
+
+        editText.setSelection(editText.editableText.indexOf("2"))
+        editText.lineBlockFormatter.applyHorizontalRule(false)
+
+        Assert.assertEquals("<ul><li>item 1</li><li>item 2</li></ul><hr /><p>test</p>", editText.toHtml())
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun addImageAfterHeadline() {
+        editText.fromHtml("<h1>Headline 1</h1>")
+
+        editText.setSelection(editText.editableText.indexOf("1"))
+        val attributes = AztecAttributes()
+        attributes.setValue("id", "1234")
+        editText.insertImage(null, attributes)
+
+        Assert.assertEquals("<h1>Headline 1</h1><img id=\"1234\" />", editText.toHtml())
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun addHRAfterHeadline() {
+        editText.fromHtml("<h1>Headline 1</h1>")
+
+        editText.setSelection(editText.editableText.indexOf("1"))
+        editText.lineBlockFormatter.applyHorizontalRule(false)
+
+        Assert.assertEquals("<h1>Headline 1</h1><hr />", editText.toHtml())
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun addImageBetweenHeadlines() {
+        editText.fromHtml("<h1>Headline 1</h1><h2>Headline 2</h2>")
+
+        editText.setSelection(editText.editableText.indexOf("1"))
+        val attributes = AztecAttributes()
+        attributes.setValue("id", "1234")
+        editText.insertImage(null, attributes)
+
+        Assert.assertEquals("<h1>Headline 1</h1><img id=\"1234\" /><h2>Headline 2</h2>", editText.toHtml())
+    }
+
+    @Test
+    @Throws(Exception::class)
+    fun addHRBetweenHeadlines() {
+        editText.fromHtml("<h1>Headline 1</h1><h2>Headline 2</h2>")
+
+        editText.setSelection(editText.editableText.indexOf("1"))
+        editText.lineBlockFormatter.applyHorizontalRule(false)
+
+        Assert.assertEquals("<h1>Headline 1</h1><hr /><h2>Headline 2</h2>", editText.toHtml())
+    }
+}