diff --git a/akka-docs/src/test/scala/docs/stream/GraphStageDocSpec.scala b/akka-docs/src/test/scala/docs/stream/GraphStageDocSpec.scala
index b5e247eac5..7be3e21a89 100644
--- a/akka-docs/src/test/scala/docs/stream/GraphStageDocSpec.scala
+++ b/akka-docs/src/test/scala/docs/stream/GraphStageDocSpec.scala
@@ -40,9 +40,11 @@ class GraphStageDocSpec extends AkkaSpec {
 
   "Demonstrate creation of GraphStage Source" in {
     //#custom-source-example
+    import akka.stream.Attributes
+    import akka.stream.Outlet
     import akka.stream.SourceShape
-    import akka.stream.Graph
     import akka.stream.stage.GraphStage
+    import akka.stream.stage.GraphStageLogic
     import akka.stream.stage.OutHandler
 
     class NumbersSource extends GraphStage[SourceShape[Int]] {