Fixed tests again

plume-oss · Jul 3, 2024 · 365c482 · 365c482
1 parent 9d33f11
commit 365c482
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 6 deletions.
diff --git a/drivers/base/src/test/scala/com/github/plume/oss/testfixtures/PlumeDriverFixture.scala b/drivers/base/src/test/scala/com/github/plume/oss/testfixtures/PlumeDriverFixture.scala
@@ -13,6 +13,7 @@ import overflowdb.{BatchedUpdate, DetachedNodeGeneric}
 
 import scala.jdk.CollectionConverters.IteratorHasAsScala
 import scala.language.postfixOps
+import scala.util.Try
 
 class PlumeDriverFixture(val driver: IDriver)
     extends AnyWordSpec
@@ -69,15 +70,25 @@ class PlumeDriverFixture(val driver: IDriver)
       val srcNode = changes
         .collectFirst {
           case c: DetachedNodeGeneric
-              if c.getRefOrId.asInstanceOf[StoredNode].id() == m.getOrElse("id", -1L).toString.toLong =>
+              if c.getRefOrId == m.getOrElse("id", -1L).toString.toLong || Try(
+                c.getRefOrId
+                  .asInstanceOf[StoredNode]
+                  .id()
+              ).map(_ == m.getOrElse("id", -1L).toString.toLong).getOrElse(false) =>
             c
         } match {
         case Some(src) => src
         case None      => fail("Unable to extract method node")
       }
       val dstNode = changes
         .collectFirst {
-          case c: NewBlock if c.getRefOrId().asInstanceOf[StoredNode].id() == b.getOrElse("id", -1L).toString.toLong =>
+          case c: NewBlock
+              if Try(c.getRefOrId().asInstanceOf[Long])
+                .map(_ == b.getOrElse("id", -1L).toString.toLong)
+                .getOrElse(false) ||
+                Try(c.getRefOrId().asInstanceOf[StoredNode].id())
+                  .map(_ == b.getOrElse("id", -1L).toString.toLong)
+                  .getOrElse(false) =>
             c
         } match {
         case Some(dst) => dst

diff --git a/drivers/overflowdb/src/test/scala/com/github/plume/oss/drivers/OverflowDbTests.scala b/drivers/overflowdb/src/test/scala/com/github/plume/oss/drivers/OverflowDbTests.scala
@@ -1,14 +1,12 @@
 package com.github.plume.oss.drivers
 
-import better.files.File
 import com.github.plume.oss.testfixtures.PlumeDriverFixture
 import com.github.plume.oss.testfixtures.PlumeDriverFixture.{b1, m1}
-import io.shiftleft.codepropertygraph.generated.{Cpg, EdgeTypes}
-import io.shiftleft.passes.IntervalKeyPool
+import io.shiftleft.codepropertygraph.generated.EdgeTypes
 import org.apache.tinkerpop.gremlin.tinkergraph.structure.TinkerGraph
 import overflowdb.BatchedUpdate
 
-import java.io.{File => JFile}
+import java.io.File as JFile
 import java.nio.charset.StandardCharsets
 import java.nio.file.{Files, Paths}
 import scala.util.{Failure, Try}