diff --git a/DCHdigi/test/test_DCHdigi/check_DCHdigi_output.py b/DCHdigi/test/test_DCHdigi/check_DCHdigi_output.py
index afc58de..b84d219 100644
--- a/DCHdigi/test/test_DCHdigi/check_DCHdigi_output.py
+++ b/DCHdigi/test/test_DCHdigi/check_DCHdigi_output.py
@@ -7,28 +7,35 @@
 #  2 : problem with calculation of hit-projection position onto the wire
 
 import ROOT
+import sys
 
-exit_code=0
-# open debug output file generated by DCHdigi alg
-f=ROOT.TFile("dch_digi_alg_debug.root")
+def main():
+    exit_code = 0
+    
+    # Open debug output file generated by DCHdigi alg
+    f = ROOT.TFile("dch_digi_alg_debug.root")
+    
+    # Retrieve the hit-wire distance distribution
+    hDpw = f.Get("hDpw")
+    hDpw.Rebin(10)
+    
+    # Retrieve the distance at which the distance distribution has its maximum
+    distance_hit_wire_more_frequent = hDpw.GetXaxis().GetBinCenter(hDpw.GetMaximumBin())
+    
+    # Check if the distance hit wire has the maximum around d=0.66cm
+    if 0.05 < abs(distance_hit_wire_more_frequent - 0.65):
+        exit_code += 1
+    
+    # Retrieve the hit-projection onto the wire to the wire distance distribution
+    hDww = f.Get("hDww")
+    
+    # Check the integral, excluding the first bin
+    hDww_integral = hDww.Integral(2, -1)
+    if hDww_integral != 0:
+        exit_code += 2
 
-# retrieve the hit-wire distance distribution
-hDpw=f.Get("hDpw")
-hDpw.Rebin(10)
-# retrieve the distance at which the distance distribution has its maximum
-distance_hit_wire_more_frequent=hDpw.GetXaxis().GetBinCenter( hDpw.GetMaximumBin() )
-# the distance hit wire has the maximum around d=0.66cm
-# if it is not the case, that means something weird is going on...
-if 0.05 < abs(distance_hit_wire_more_frequent - 0.65) :
-    exit_code+=1
+    return exit_code
 
-# retrieve the hit-projection onto the wire to the wire distance distribution
-hDww=f.Get("hDww")
-# since the hit-projection onto the wire should be a point on the wire, the distance should be zero
-# and all the counts are pushed to the bin number 1 (which is excluded from the integral)
-hDww_integral=hDww.Integral(2,-1)
-if hDww_integral != 0 :
-    exit_code+=2
-
-# we have to print the exit code, so it can be captured by the bash script
-print(exit_code)
+if __name__ == "__main__":
+    code = main()
+    sys.exit(code)
diff --git a/DCHdigi/test/test_DCHdigi/test_DCHdigi.sh b/DCHdigi/test/test_DCHdigi/test_DCHdigi.sh
index f218ccd..6163b6e 100644
--- a/DCHdigi/test/test_DCHdigi/test_DCHdigi.sh
+++ b/DCHdigi/test/test_DCHdigi/test_DCHdigi.sh
@@ -21,7 +21,4 @@ fi
 k4run runDCHdigi.py
 
 # check distribution of distance from hit position to the wire
-check_DCHdigi_output=$( (python3 check_DCHdigi_output.py) 2>&1)
-
-# return value printed out by the previous python script
-exit $check_DCHdigi_output
+python3 check_DCHdigi_output.py