diff --git a/ImportGraph/Cli.lean b/ImportGraph/Cli.lean
index 0c126bb..3bd3c95 100644
--- a/ImportGraph/Cli.lean
+++ b/ImportGraph/Cli.lean
@@ -8,6 +8,7 @@ import Batteries.Lean.IO.Process
 import ImportGraph.CurrentModule
 import ImportGraph.Imports
 import ImportGraph.Lean.Name
+import ImportGraph.Unused
 
 open Cli
 
@@ -17,15 +18,15 @@ open ImportGraph
 /-- Write an import graph, represented as a `NameMap (Array Name)` to the ".dot" graph format.
     Nodes in the `unused` set will be shaded light gray. -/
 def asDotGraph
-    (graph : NameMap (Array Name)) (unused : NameSet := {}) (header := "import_graph") :
+    (graph : NameMap (Array Name)) (used : NameSet := {}) (header := "import_graph") :
     String := Id.run do
   let mut lines := #[s!"digraph \"{header}\" " ++ "{"]
-  -- Add node styling for unused nodes
-  lines := lines.push "  node [style=filled];"
-  -- Add nodes and edges
   for (n, is) in graph do
-    if unused.contains n then
-      lines := lines.push s!"  \"{n}\" [fillcolor=lightgray];"
+    if used.contains n then
+      lines := lines.push s!"  \"{n}\";"
+    else
+      lines := lines.push s!"  \"{n}\" [style=filled, fillcolor=lightgray];"
+    -- Then add edges
     for i in is do
       lines := lines.push s!"  \"{i}\" -> \"{n}\";"
   lines := lines.push "}"
@@ -45,10 +46,10 @@ def importGraphCLI (args : Cli.Parsed) : IO UInt32 := do
   searchPathRef.set compile_time_search_path%
   let dotFile ← try unsafe withImportModules #[{module := to}] {} (trustLevel := 1024) fun env => do
     let mut graph := env.importGraph
-    let ctx := { options := {}, fileName := "<input>", fileMap := default }
-    let state := { env }
-    let used ← Prod.fst <$> (CoreM.toIO (env.transitivelyRequiredModules to) ctx state)
-    let unused := graph.fold (fun acc n _ => if used.contains n then acc else acc.insert n) NameSet.empty
+    let used ← do
+      let ctx := { options := {}, fileName := "<input>", fileMap := default }
+      let state := { env }
+      Prod.fst <$> (CoreM.toIO (env.transitivelyRequiredModules to) ctx state)
     if let Option.some f := from? then
       graph := graph.downstreamOf (NameSet.empty.insert f)
     if ¬(args.hasFlag "include-deps") then
@@ -65,7 +66,7 @@ def importGraphCLI (args : Cli.Parsed) : IO UInt32 := do
       graph := graph.filterGraph filterMathlibMeta (replacement := `«Mathlib.Tactics»)
     if args.hasFlag "reduce" then
       graph := graph.transitiveReduction
-    return asDotGraph graph (unused := unused)
+    return asDotGraph graph (used := used)
   catch err =>
     -- TODO: try to build `to` first, so this doesn't happen
     throw <| IO.userError <| s!"{err}\nIf the error above says `unknown package`, " ++
diff --git a/ImportGraph/RequiredModules.lean b/ImportGraph/RequiredModules.lean
index 51f04e9..be67431 100644
--- a/ImportGraph/RequiredModules.lean
+++ b/ImportGraph/RequiredModules.lean
@@ -72,6 +72,7 @@ and identifies all modules contains constants which are transitively required by
 def Environment.transitivelyRequiredModules (env : Environment) (module : Name) : CoreM NameSet := do
   let mut requiredModules : NameSet := {}
   for (_, ci) in env.constants.map₁.toList do
+    if !ci.name.isInternal then
     if let some m := env.getModuleFor? ci.name then
       if m = module then
         for r in (← ci.name.transitivelyRequiredModules env) do
diff --git a/ImportGraph/Unused.lean b/ImportGraph/Unused.lean
new file mode 100644
index 0000000..224f00d
--- /dev/null
+++ b/ImportGraph/Unused.lean
@@ -0,0 +1,3 @@
+/-!
+This is a dummy file to test the `unused` flag in the import graph.
+-/
diff --git a/ImportGraphTest.lean b/ImportGraphTest.lean
new file mode 100644
index 0000000..eb20b90
--- /dev/null
+++ b/ImportGraphTest.lean
@@ -0,0 +1,2 @@
+import ImportGraphTest.Imports
+import ImportGraphTest.Dot
diff --git a/ImportGraphTest/.gitignore b/ImportGraphTest/.gitignore
new file mode 100644
index 0000000..ed91303
--- /dev/null
+++ b/ImportGraphTest/.gitignore
@@ -0,0 +1 @@
+produced.dot
\ No newline at end of file
diff --git a/ImportGraphTest/Dot.lean b/ImportGraphTest/Dot.lean
new file mode 100644
index 0000000..ddb203b
--- /dev/null
+++ b/ImportGraphTest/Dot.lean
@@ -0,0 +1,29 @@
+import ImportGraph.Cli
+
+def readFile (path : System.FilePath) : IO String :=
+  IO.FS.readFile path
+
+def runGraphCommand : IO Unit := do
+  let _ ← IO.Process.output {
+    cmd := "lake"
+    args := #["exe", "graph", "--to", "ImportGraph.Cli", "ImportGraphTest/produced.dot"]
+  }
+
+def compareOutputs (expected : String) (actual : String) : IO Bool := do
+  let expectedLines := expected.splitOn "\n" |>.filter (·.trim.length > 0) |>.map String.trim
+  let actualLines := actual.splitOn "\n" |>.filter (·.trim.length > 0) |>.map String.trim
+  pure (expectedLines == actualLines)
+
+/-- info: Test passed: The graph command output matches the expected.dot file. -/
+#guard_msgs in
+#eval show IO Unit from do
+  runGraphCommand
+  let expectedOutput ← readFile "ImportGraphTest/expected.dot"
+  let actualOutput ← readFile "ImportGraphTest/produced.dot"
+  let isEqual ← compareOutputs expectedOutput actualOutput
+  if isEqual then
+    IO.println "Test passed: The graph command output matches the expected.dot file."
+  else
+    IO.println "Test failed: The graph command output does not match the expected.dot file."
+    IO.println s!"Expected:\n{expectedOutput}"
+    IO.println s!"Actual:\n{actualOutput}"
diff --git a/test/Imports.lean b/ImportGraphTest/Imports.lean
similarity index 81%
rename from test/Imports.lean
rename to ImportGraphTest/Imports.lean
index fa733bf..6688366 100644
--- a/test/Imports.lean
+++ b/ImportGraphTest/Imports.lean
@@ -22,7 +22,17 @@ ImportGraph.RequiredModules
 #guard_msgs in
 #find_home importTest
 
-open Elab Command in
+open Elab Command
+
+elab "#transitivelyRequiredModules_test" : command => do
+  let env ← getEnv
+  let unused ← liftCoreM <| env.transitivelyRequiredModules `ImportGraph.RequiredModules
+  logInfo s!"{unused.contains `Init.Data.Option.Lemmas}"
+
+/-- info: true -/
+#guard_msgs in
+#transitivelyRequiredModules_test
+
 elab "#my_test" : command => do
   -- functionality of `#redundant_imports`
   let expected := #[`ImportGraph.RequiredModules]
diff --git a/ImportGraphTest/expected.dot b/ImportGraphTest/expected.dot
new file mode 100644
index 0000000..7a6af16
--- /dev/null
+++ b/ImportGraphTest/expected.dot
@@ -0,0 +1,13 @@
+digraph "import_graph" {
+  "ImportGraph.Unused" [style=filled, fillcolor=lightgray];
+  "ImportGraph.CurrentModule";
+  "ImportGraph.Lean.Name";
+  "ImportGraph.RequiredModules";
+  "ImportGraph.Imports";
+  "ImportGraph.RequiredModules" -> "ImportGraph.Imports";
+  "ImportGraph.Cli";
+  "ImportGraph.CurrentModule" -> "ImportGraph.Cli";
+  "ImportGraph.Imports" -> "ImportGraph.Cli";
+  "ImportGraph.Lean.Name" -> "ImportGraph.Cli";
+  "ImportGraph.Unused" -> "ImportGraph.Cli";
+}
\ No newline at end of file
diff --git a/lakefile.toml b/lakefile.toml
index f8404dc..837cf62 100644
--- a/lakefile.toml
+++ b/lakefile.toml
@@ -1,6 +1,6 @@
 name = "importGraph"
 defaultTargets = ["ImportGraph", "graph"]
-testRunner = "test"
+testRunner = "ImportGraphTest"
 
 [[require]]
 name = "Cli"
@@ -24,6 +24,5 @@ root = "Main"
 # Remove this line if you do not need such functionality.
 supportInterpreter = true
 
-[[lean_exe]]
-name = "test"
-srcDir = "scripts"
+[[lean_lib]]
+name = "ImportGraphTest"
diff --git a/scripts/test.lean b/scripts/test.lean
deleted file mode 100644
index fcc6f5f..0000000
--- a/scripts/test.lean
+++ /dev/null
@@ -1,53 +0,0 @@
-/-
-Copyright (c) 2024 Kim Morrison. All rights reserved.
-Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Kim Morrison
--/
-open IO.Process
-open System
-
-/-
-This is a copy-paste from Batteries test runner.
-If https://github.com/leanprover/lean4/pull/4121 is resolved, use Batteries' script directly.
--/
-
-/--
-Run tests.
-
-If no arguments are provided, run everything in the `test/` directory.
-
-If arguments are provided, assume they are names of files in the `test/` directory,
-and just run those.
--/
-def main (args : List String) : IO Unit := do
-  -- We first run `lake build`, to ensure oleans are available.
-  -- Alternatively, we could use `lake lean` below instead of `lake env lean`,
-  -- but currently with parallelism this results in build jobs interfering with each other.
-  _ ← (← IO.Process.spawn { cmd := "lake", args := #["build"] }).wait
-  -- Collect test targets, either from the command line or by walking the `test/` directory.
-  let targets ← match args with
-  | [] => System.FilePath.walkDir "./test"
-  | _ => pure <| (args.map fun t => mkFilePath [".", "test", t] |>.withExtension "lean") |>.toArray
-  let existing ← targets.filterM fun t => do pure <| (← t.pathExists) && !(← t.isDir)
-  -- Generate a `lake env lean` task for each test target.
-  let tasks ← existing.mapM fun t => do
-    IO.asTask do
-      let out ← IO.Process.output
-        { cmd := "lake",
-          args := #["env", "lean", t.toString],
-          env := #[("LEAN_ABORT_ON_PANIC", "1")] }
-      if out.exitCode = 0 then
-        IO.println s!"Test succeeded: {t}"
-      else
-        IO.eprintln s!"Test failed: `lake env lean {t}` produced:"
-        unless out.stdout.isEmpty do IO.eprintln out.stdout
-        unless out.stderr.isEmpty do IO.eprintln out.stderr
-      pure out.exitCode
-  -- Wait on all the jobs and exit with 1 if any failed.
-  let mut exitCode : UInt8 := 0
-  for t in tasks do
-    let e ← IO.wait t
-    match e with
-    | .ok 0 => pure ()
-    | _ => exitCode := 1
-  exit exitCode