diff --git a/aiida_worktree/__init__.py b/aiida_worktree/__init__.py
index 5665dd81..8442e43c 100644
--- a/aiida_worktree/__init__.py
+++ b/aiida_worktree/__init__.py
@@ -2,4 +2,7 @@
 from .node import Node
 from .decorator import node, build_node
 
+
 __version__ = "0.0.4"
+
+__all__ = ["WorkTree", "Node", "node", "build_node"]
diff --git a/aiida_worktree/cli/__init__.py b/aiida_worktree/cli/__init__.py
index 38730244..8de3c0ed 100644
--- a/aiida_worktree/cli/__init__.py
+++ b/aiida_worktree/cli/__init__.py
@@ -3,3 +3,6 @@
 The commands need to be imported here for them to be registered with the top-level command group.
 """
 from aiida_worktree.cli import cmd_tree
+
+
+__all__ = ["cmd_tree"]
diff --git a/aiida_worktree/engine/worktree.py b/aiida_worktree/engine/worktree.py
index d8366a76..2114658a 100644
--- a/aiida_worktree/engine/worktree.py
+++ b/aiida_worktree/engine/worktree.py
@@ -19,7 +19,7 @@
 
 
 from aiida.engine.processes.exit_code import ExitCode
-from aiida.engine.processes.process import Process, ProcessState
+from aiida.engine.processes.process import Process
 
 from aiida.engine.processes.workchains.awaitable import (
     Awaitable,
@@ -30,9 +30,6 @@
 from aiida.engine.processes.workchains.workchain import Protect, WorkChainSpec
 from aiida.engine import run_get_node
 
-from os.path import splitext
-import yaml
-
 
 if t.TYPE_CHECKING:
     from aiida.engine.runners import Runner  # pylint: disable=unused-import
@@ -275,8 +272,10 @@ def _do_step(self) -> t.Any:
 
         # If the worktree is finished or the result is an ExitCode, we exit by returning
         if finished:
-            result = self.finalize()
-            return result
+            if isinstance(result, ExitCode):
+                return result
+            else:
+                return self.finalize()
 
         if self._awaitables:
             return Wait(self._do_step, "Waiting before next step")
@@ -684,8 +683,6 @@ def run_nodes(self, names):
                 self.to_context(**{name: process})
             elif node["metadata"]["node_type"] in ["worktree"]:
                 # process = run_get_node(executor, *args, **kwargs)
-                from aiida_worktree.utils import merge_properties
-                from aiida.orm.utils.serialize import serialize
 
                 print("node  type: worktree.")
                 wt = self.run_executor(executor, args, kwargs, var_args, var_kwargs)
@@ -859,9 +856,6 @@ def check_parent_state(self, name):
                         "SKIPPED",
                         "FAILED",
                     ]:
-                        # print(
-                        #     f"    {name}: Input node {link['from_node']}, {self.ctx.nodes[link['from_node']]['state']} ."
-                        # )
                         ready = False
                         return (
                             ready,
diff --git a/aiida_worktree/executors/qe.py b/aiida_worktree/executors/qe.py
index c8942e85..3b087f2f 100644
--- a/aiida_worktree/executors/qe.py
+++ b/aiida_worktree/executors/qe.py
@@ -1,6 +1,4 @@
-import time
-from aiida import engine, orm, load_profile
-from aiida.orm import load_code
+from aiida import load_profile
 
 load_profile()
 
diff --git a/aiida_worktree/nodes/builtin.py b/aiida_worktree/nodes/builtin.py
index dc8dd495..4ee2888f 100644
--- a/aiida_worktree/nodes/builtin.py
+++ b/aiida_worktree/nodes/builtin.py
@@ -1,5 +1,4 @@
 from aiida_worktree.node import Node
-from aiida_worktree.executors.builtin import GatherWorkChain
 
 
 class AiiDAGather(Node):
diff --git a/aiida_worktree/nodes/qe.py b/aiida_worktree/nodes/qe.py
index 203f5be6..388b576e 100644
--- a/aiida_worktree/nodes/qe.py
+++ b/aiida_worktree/nodes/qe.py
@@ -1,5 +1,4 @@
 from aiida_worktree.node import Node
-from aiida import orm
 
 
 class AiiDAKpoint(Node):
diff --git a/aiida_worktree/worktree.py b/aiida_worktree/worktree.py
index 38b50183..168686e4 100644
--- a/aiida_worktree/worktree.py
+++ b/aiida_worktree/worktree.py
@@ -43,7 +43,6 @@ def run(self):
         """
         from aiida_worktree.engine.worktree import WorkTree as WorkTreeEngine
         from aiida_worktree.utils import merge_properties
-        from aiida.orm.utils.serialize import serialize
         from aiida.manage import manager
 
         # One can not run again if the process is alreay created. otherwise, a new process node will
@@ -253,7 +252,6 @@ def play_nodes(self, nodes):
 
     def reset(self):
         """Reset the worktree."""
-        from aiida.engine.processes import control
 
         self.process = None
         for node in self.nodes:
diff --git a/tests/test_calcjob.py b/tests/test_calcjob.py
index 08f61964..cb4aecac 100644
--- a/tests/test_calcjob.py
+++ b/tests/test_calcjob.py
@@ -1,5 +1,3 @@
-import time
-import numpy as np
 import aiida
 
 aiida.load_profile()