Fix plusType for schema that contains only ISL 2.0 version marker

amazon-ion · Feb 20, 2024 · ea13fba · ea13fba
1 parent bbab678
commit ea13fba
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 3 deletions.
diff --git a/ion-schema/src/main/kotlin/com/amazon/ionschema/internal/SchemaImpl_2_0.kt b/ion-schema/src/main/kotlin/com/amazon/ionschema/internal/SchemaImpl_2_0.kt
@@ -327,14 +327,15 @@ internal class SchemaImpl_2_0 internal constructor(
 
         val newIsl = mutableListOf<IonValue>()
         var newTypeAdded = false
-        isl.forEachIndexed { idx, value ->
+        isl.forEach { value ->
             if (!newTypeAdded) {
                 if (isType(value) && (value["name"] as? IonSymbol)?.stringValue() == type.name) {
                     // new type replaces existing type of the same name
                     newIsl.add(type.isl.clone())
                     newTypeAdded = true
-                    return@forEachIndexed
-                } else if (value.hasTypeAnnotation("schema_footer") || idx == isl.lastIndex) {
+                    return@forEach
+                } else if (value.hasTypeAnnotation("schema_footer")) {
+                    // Insert the new type right before the footer
                     newIsl.add(type.isl.clone())
                     newTypeAdded = true
                 }

diff --git a/ion-schema/src/test/kotlin/com/amazon/ionschema/internal/SchemaImpl_2_0_Test.kt b/ion-schema/src/test/kotlin/com/amazon/ionschema/internal/SchemaImpl_2_0_Test.kt
@@ -2,6 +2,7 @@ package com.amazon.ionschema.internal
 
 import com.amazon.ion.IonStruct
 import com.amazon.ion.IonValue
+import com.amazon.ionschema.ION
 import com.amazon.ionschema.IonSchemaSystemBuilder
 import com.amazon.ionschema.IonSchemaTests
 import com.amazon.ionschema.IonSchemaVersion
@@ -65,6 +66,67 @@ class SchemaImpl_2_0_Test {
         assertEquals(expectedNewFooTypeIsl, newFooTypeIsl)
     }
 
+    @Test
+    fun `plusType(type) should correctly insert new type when a schema footer is present`() {
+        val schema = ISS.newSchema("""
+            ${'$'}ion_schema_2_0
+            schema_header::{}
+            type::{ name: foo }
+            schema_footer::{}
+        """)
+        val newType = schema.newType("type::{ name: bar }")
+        val newSchema = schema.plusType(newType)
+
+        val oldSchemaTypes = schema.getDeclaredTypes().asSequence().map { it.name }.toSet()
+        assertEquals(setOf("foo"), oldSchemaTypes, "The original schema should not be modified.")
+
+        val newSchemaTypes = newSchema.getDeclaredTypes().asSequence().map { it.name }.toSet()
+        val expectedNewSchemaTypes = setOf("foo", "bar")
+        assertEquals(expectedNewSchemaTypes, newSchemaTypes)
+    }
+
+    @Test
+    fun `plusType(type) should correctly insert a new type into a schema that is empty aside from its version marker`() {
+        val schema = ISS.newSchema("\$ion_schema_2_0")
+        val newType = schema.newType("type::{ name: bar }")
+        val newSchema = schema.plusType(newType)
+
+        val oldSchemaTypes = schema.getDeclaredTypes().asSequence().map { it.name }.toSet()
+        assertEquals(emptySet<Nothing>(), oldSchemaTypes, "The original schema should not be modified.")
+
+        val newSchemaTypes = newSchema.getDeclaredTypes().asSequence().map { it.name }.toSet()
+        val expectedNewSchemaTypes = setOf("bar")
+        assertEquals(expectedNewSchemaTypes, newSchemaTypes)
+    }
+
+    @Test
+    fun `plusType(type) should preserve top-level open content in a schema`() {
+        val schema = ISS.newSchema("""
+            ${'$'}ion_schema_2_0
+            abc
+            type::{ name: foo }
+            def
+        """)
+        val newType = schema.newType("type::{ name: bar }")
+        val newSchema = schema.plusType(newType)
+
+        val oldSchemaTypes = schema.getDeclaredTypes().asSequence().map { it.name }.toSet()
+        assertEquals(setOf("foo"), oldSchemaTypes, "The original schema should not be modified.")
+
+        val newSchemaTypes = newSchema.getDeclaredTypes().asSequence().map { it.name }.toSet()
+        val expectedNewSchemaTypes = setOf("foo", "bar")
+        assertEquals(expectedNewSchemaTypes, newSchemaTypes)
+
+        val expectedNewSchemaIsl = ION.loader.load("""
+            ${'$'}ion_schema_2_0
+            abc
+            type::{ name: foo }
+            def
+            type::{ name: bar }
+        """)
+        assertEquals(newSchema.isl, expectedNewSchemaIsl)
+    }
+
     @Test
     fun `getType(name) should return a declared type`() {
         val schema = ISS.newSchema("\$ion_schema_2_0 type::{ name: foo }")