diff --git a/mockgen.go b/mockgen.go
index 626a4cc4776..5b7cd4f0a3c 100644
--- a/mockgen.go
+++ b/mockgen.go
@@ -17,6 +17,3 @@ package quic
 //go:generate sh -c "./mockgen_private.sh quic mock_unknown_packet_handler_test.go github.com/lucas-clemente/quic-go unknownPacketHandler"
 //go:generate sh -c "./mockgen_private.sh quic mock_packet_handler_manager_test.go github.com/lucas-clemente/quic-go packetHandlerManager"
 //go:generate sh -c "./mockgen_private.sh quic mock_multiplexer_test.go github.com/lucas-clemente/quic-go multiplexer"
-
-//go:generate sh -c "find . -type f -name 'mock_*_test.go' | xargs sed -i '' 's/quic_go.//g'"
-//go:generate sh -c "goimports -w mock*_test.go"
diff --git a/mockgen_private.sh b/mockgen_private.sh
index 07b08d940d8..0ba5f64ef81 100755
--- a/mockgen_private.sh
+++ b/mockgen_private.sh
@@ -21,4 +21,8 @@ export GOPATH="$TEMP_DIR:$GOPATH"
 
 mockgen -package $1 -self_package $1 -destination $2 $3 $INTERFACE_NAME
 
+# mockgen imports quic-go as 'import quic_go github.com/lucas_clemente/quic-go'
+sed -i '' 's/quic_go.//g' $2
+goimports -w $2
+
 rm -r "$TEMP_DIR"