diff --git a/aiida_aimall/calculations.py b/aiida_aimall/calculations.py
index d759c4d..77ba3a9 100644
--- a/aiida_aimall/calculations.py
+++ b/aiida_aimall/calculations.py
@@ -75,10 +75,10 @@ def define(cls, spec):
     # ---------------------------------------------------
 
     def prepare_for_submission(self, folder):
-        """
-        Create input files.
-        :param folder: an `aiida.common.folders.Folder` where the plugin should temporarily place all files
-            needed by the calculation.
+        """Create input files.
+
+        :param folder: an `aiida.common.folders.Folder` where the plugin should temporarily
+            place all files needed by the calculation.
         :return: `aiida.common.datastructures.CalcInfo` instance
         """
         input_string = self.inputs.file.get_content()
diff --git a/examples/example_01.py b/examples/example_01.py
index bc1a0ea..4cf52b1 100644
--- a/examples/example_01.py
+++ b/examples/example_01.py
@@ -4,15 +4,16 @@
 Usage: ./example_01.py
 """
 from os import path
+import sys
 
 import click
-import sys
+
+from aiida import engine
 from aiida.common import NotExistent
-from aiida import cmdline, engine
 from aiida.orm import Code
 from aiida.plugins import CalculationFactory, DataFactory
 
-from aiida_aimall import helpers
+# from aiida_aimall import helpers
 
 INPUT_DIR = path.join(path.dirname(path.realpath(__file__)), "input_files")
 
@@ -22,10 +23,10 @@ def test_run(aimall_code):
 
     Uses test helpers to create AiiDA Code on the fly.
     """
-    if not aimall_code:
-        # get code
-        computer = helpers.get_computer()
-        aimall_code = helpers.get_code(entry_point="aimall", computer=computer)
+    # if not aimall_code:
+    #     # get code
+    #     computer = helpers.get_computer()
+    #     aimall_code = helpers.get_code(entry_point="aimall", computer=computer)
 
     # Prepare input parameters
     AimqbParameters = DataFactory("aimall")
@@ -33,7 +34,6 @@ def test_run(aimall_code):
 
     SinglefileData = DataFactory("singlefile")
     file = SinglefileData(file=path.join(INPUT_DIR, "file1.txt"))
-    
 
     # set up calculation
     inputs = {
@@ -49,13 +49,14 @@ def test_run(aimall_code):
     # from aiida.engine import submit
     # future = submit(CalculationFactory('aimall'), **inputs)
     result = engine.run(CalculationFactory("aimall"), **inputs)
-
+    at_result = list(result["atomic_properties"].get_dict().keys())[0]
+    print(f"Calculation ran, first atom is {at_result}")
     # computed_aimqb = result["aimall"].get_content()
     # print(f"Computed aimqb: \n{computed_aimqb}")
 
 
 @click.command()
-@click.argument("codelabel",default="aimqb@localhost")
+@click.argument("codelabel", default="aimall2@localhost")
 # @cmdline.utils.decorators.with_dbenv()
 # @cmdline.params.options.CODE()
 def cli(codelabel):
@@ -69,8 +70,10 @@ def cli(codelabel):
     """
     try:
         code = Code.get_from_string(codelabel)
+        if code:
+            print(f"The code {codelabel} does exist")
     except NotExistent:
-        print(f'The code {codelabel} does not exist')
+        print(f"The code {codelabel} does not exist")
         sys.exit(1)
     test_run(codelabel)