chore: make sure nodes support valid components

mongodb · Jul 8, 2024 · bc689c5 · bc689c5
1 parent cc99300
commit bc689c5
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 5 deletions.
diff --git a/packages/mongodb-mql-model/src/main/kotlin/com/mongodb/jbplugin/mql/Node.kt b/packages/mongodb-mql-model/src/main/kotlin/com/mongodb/jbplugin/mql/Node.kt
@@ -37,6 +37,8 @@ data class Node<S>(
 ) {
     inline fun <reified C : Component> component(): C? = components.firstOrNull { it is C } as C?
 
+    fun <C : Component> component(withClass: Class<C>): C? = components.firstOrNull { withClass.isInstance(it) } as C?
+
     inline fun <reified C : Component> components(): List<C> = components.filterIsInstance<C>()
 
     inline fun <reified C : Component> hasComponent(): Boolean = component<C>() != null

diff --git a/packages/mongodb-mql-model/src/test/kotlin/com/mongodb/jbplugin/mql/NodeTest.kt b/packages/mongodb-mql-model/src/test/kotlin/com/mongodb/jbplugin/mql/NodeTest.kt
@@ -1,9 +1,10 @@
 package com.mongodb.jbplugin.mql
 
-import com.mongodb.jbplugin.mql.components.HasFieldReference
-import com.mongodb.jbplugin.mql.components.Named
+import com.mongodb.jbplugin.mql.components.*
 import org.junit.jupiter.api.Assertions.*
 import org.junit.jupiter.api.Test
+import org.junit.jupiter.params.ParameterizedTest
+import org.junit.jupiter.params.provider.MethodSource
 
 class NodeTest {
     @Test
@@ -27,9 +28,14 @@ class NodeTest {
         val node =
             Node<Unit?>(
                 null,
-                listOf(HasFieldReference(HasFieldReference.Known("field1")), HasFieldReference(HasFieldReference.Known(
-"field2"
-))),
+                listOf(
+                    HasFieldReference(HasFieldReference.Known("field1")),
+                    HasFieldReference(
+                        HasFieldReference.Known(
+                            "field2",
+                        ),
+                    ),
+                ),
             )
         val fieldReferences = node.components<HasFieldReference>()
 
@@ -60,4 +66,42 @@ class NodeTest {
 
         assertFalse(hasNamedComponent)
     }
+
+    @MethodSource("validComponents")
+    @ParameterizedTest
+    fun `does support the following component`(
+        component: Component,
+        componentClass: Class<Component>,
+    ) {
+        val node =
+            Node<Unit?>(
+                null,
+                listOf(
+                    component,
+                ),
+            )
+
+        assertNotNull(node.component(componentClass))
+    }
+
+    companion object {
+        @JvmStatic
+        fun validComponents(): Array<Array<Any>> =
+            arrayOf(
+                arrayOf(HasChildren<Unit?>(emptyList()), HasChildren::class.java),
+                arrayOf(HasCollectionReference(HasCollectionReference.Unknown), HasCollectionReference::class.java),
+                arrayOf(HasCollectionReference(HasCollectionReference.Known(Namespace("db", "coll"))),
+ HasCollectionReference::class.java),
+                arrayOf(HasCollectionReference(HasCollectionReference.OnlyCollection("coll")),
+ HasCollectionReference::class.java),
+                arrayOf(HasFieldReference(HasFieldReference.Unknown), HasFieldReference::class.java),
+                arrayOf(HasFieldReference(HasFieldReference.Known("abc")), HasFieldReference::class.java),
+                arrayOf(HasFilter<Unit?>(Node(null, emptyList())), HasFilter::class.java),
+                arrayOf(HasValueReference(HasValueReference.Unknown), HasValueReference::class.java),
+                arrayOf(HasValueReference(HasValueReference.Constant(123, "int")), HasValueReference::class.java),
+                arrayOf(HasValueReference(HasValueReference.Runtime("int")), HasValueReference::class.java),
+                arrayOf(HasValueReference(HasValueReference.Unknown), HasValueReference::class.java),
+                arrayOf(Named("abc"), Named::class.java),
+            )
+    }
 }