fix: fix tests

unit-mesh · Aug 30, 2023 · f1485bf · f1485bf
1 parent b24ffb5
commit f1485bf
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 8 deletions.
diff --git a/src/main/kotlin/cc/unitmesh/cf/domains/DomainDispatcher.kt b/src/main/kotlin/cc/unitmesh/cf/domains/DomainDispatcher.kt
@@ -7,6 +7,7 @@ import cc.unitmesh.cf.infrastructure.cache.CachedEmbedding
 import cc.unitmesh.cf.infrastructure.llms.embedding.Embedding
 import org.springframework.stereotype.Component
 import org.reflections.Reflections
+import org.reflections.scanners.SubTypesScanner
 
 @Component
 class DomainDispatcher(
@@ -18,11 +19,9 @@ class DomainDispatcher(
     }
 
     fun lookupDomains(): List<Class<out DomainDetector>> {
-        val reflections = Reflections(this.javaClass.`package`.name)
-        return reflections.getTypesAnnotatedWith(Domain::class.java)
-            .filter {
-                it.superclass == DomainDetector::class.java
-            }
-            .toList() as List<Class<out DomainDetector>>
+        val reflections = Reflections(DomainDispatcher::class.java.`package`.name, SubTypesScanner(false))
+
+        return reflections.getSubTypesOf(DomainDetector::class.java)
+            .toList()
     }
 }
diff --git a/src/test/kotlin/cc/unitmesh/cf/domains/DomainDispatcherTest.kt b/src/test/kotlin/cc/unitmesh/cf/domains/DomainDispatcherTest.kt
@@ -2,6 +2,7 @@ package cc.unitmesh.cf.domains
 
 import cc.unitmesh.cf.infrastructure.cache.CachedEmbedding
 import io.kotest.matchers.ints.shouldBeGreaterThan
+import io.kotest.matchers.ints.shouldBeGreaterThanOrEqual
 import io.mockk.MockKAnnotations
 import io.mockk.impl.annotations.MockK
 import io.mockk.junit5.MockKExtension
@@ -20,11 +21,10 @@ class DomainDispatcherTest {
     fun setUp() = MockKAnnotations.init(this, relaxUnitFun = true) // turn relaxUnitFun on for all mocks
 
     @Test
-    @Disabled
     fun should_enable_lookup_domains() {
         val domainDispatcher = DomainDispatcher(cachedEmbedding)
         val domain = domainDispatcher.lookupDomains()
 
-        domain.size shouldBeGreaterThan  1
+        domain.size shouldBeGreaterThanOrEqual   1
     }
 }