diff --git a/Tests/Data/HydroMechanics/GroundEquilibrium/pythonBCsOGS.py b/Tests/Data/HydroMechanics/GroundEquilibrium/pythonBCsOGS.py
index 05097ea717f3df6c2a0cf07340af90c036db3db9..baaffb4fa0f78667e49396c1d4eb749440756ea2 100644
--- a/Tests/Data/HydroMechanics/GroundEquilibrium/pythonBCsOGS.py
+++ b/Tests/Data/HydroMechanics/GroundEquilibrium/pythonBCsOGS.py
@@ -1,21 +1,25 @@
 # Collection of python boundary condition (BC) classes for OpenGeoSys
 
-import OpenGeoSys
+try:
+    import ogs.callbacks as OpenGeoSys
+except ModuleNotFoundError:
+    import OpenGeoSys
+
 import numpy as np
 
-s_a = 365.25*24*3600 #=31557600 seconds per year
+s_a = 365.25 * 24 * 3600  # =31557600 seconds per year
 
 # Choose parametrization
-g = 9.81 #m/s²
-nu = 0.4 # -
-phi = 0.2 #-
-rho_W = 1000.0 #kg/m³
-rho_S = 3000.0 #kg/m³
-u_max = -1.0 #m
+g = 9.81  # m/s²
+nu = 0.4  # -
+phi = 0.2  # -
+rho_W = 1000.0  # kg/m³
+rho_S = 3000.0  # kg/m³
+u_max = -1.0  # m
 
-Lx=10 #m
-Ly=10 #m
-T=1e6 #s
+Lx = 10  # m
+Ly = 10  # m
+T = 1e6  # s
 
 # Nomenclature: BC Process_LocationQuantity_Component
 # 					(THM)			(XYZ)
@@ -27,87 +31,90 @@ T=1e6 #s
 
 
 def ExternalDisplacement(x, t):
-	uy = u_max*(t/T)*(x/Lx)**2
-	
-	return uy
+    uy = u_max * (t / T) * (x / Lx) ** 2
+
+    return uy
+
 
 # Hydraulic BCs
 # -------------
 class BCH_SurfacePressure(OpenGeoSys.BoundaryCondition):
+    def getDirichletBCValue(self, t, coords, node_id, primary_vars):
+        x, y, z = coords
 
-	def getDirichletBCValue(self, t, coords, node_id, primary_vars):
-		x, y, z = coords
+        value = 0.0
 
-		value = 0.0
+        return (True, value)
 
-		return (True, value)
 
 class BCH_VerticalPressure(OpenGeoSys.BoundaryCondition):
+    def getDirichletBCValue(self, t, coords, node_id, primary_vars):
+        x, y, z = coords
 
-	def getDirichletBCValue(self, t, coords, node_id, primary_vars):
-		x, y, z = coords
+        value = -rho_W * g * y
 
-		value = -rho_W * g * y
-
-		return (True, value)
+        return (True, value)
 
 
 # Mechanics BCs
 # -------------
 class BCM_VerticalTraction_X(OpenGeoSys.BoundaryCondition):
-	
-	def getFlux(self, t, coords, primary_vars): #here Neumann BC: flux of linear momentum
-		x, y, z = coords
-		
-		value = (nu/(1-nu) * ((phi-1)*rho_W + (1-phi)*rho_S) + rho_W) * g*y
-		jacob = [0.0, 0.0, 0.0]
-		
-		return (True, value, jacob)
+    def getFlux(
+        self, t, coords, primary_vars
+    ):  # here Neumann BC: flux of linear momentum
+        x, y, z = coords
+
+        value = (
+            (nu / (1 - nu) * ((phi - 1) * rho_W + (1 - phi) * rho_S) + rho_W) * g * y
+        )
+        jacob = [0.0, 0.0, 0.0]
+
+        return (True, value, jacob)
+
 
 class BCM_MonitoringTraction(OpenGeoSys.BoundaryCondition):
+    def getFlux(self, t, coords, primary_vars):
+        x, y, z = coords
+        pp_actual = primary_vars[0]
+        ux_actual = primary_vars[1]
+        uy_actual = primary_vars[2]
+
+        uy_target = ExternalDisplacement(x, t)
+        uy_diff = uy_actual - uy_target
+
+        if abs(uy_diff) > 1e-15:
+            print("WARNING: uy_diff = ", uy_diff)
+
+        value = 0.0
+        jacob = [0.0, 0.0, 0.0]
+
+        return (False, value, jacob)
 
-	def getFlux(self, t, coords, primary_vars):
-		x, y, z = coords
-		pp_actual = primary_vars[0]
-		ux_actual = primary_vars[1]
-		uy_actual = primary_vars[2]
-		
-		uy_target = ExternalDisplacement(x, t)
-		uy_diff = uy_actual-uy_target
-		
-		if (abs(uy_diff) > 1e-15):
-			print("WARNING: uy_diff = ", uy_diff)
-		
-		value = 0.0
-		jacob = [0.0, 0.0, 0.0]
-		
-		return (False, value, jacob)
 
 class BCM_SurfaceDisplacement_Y(OpenGeoSys.BoundaryCondition):
-	
-	def getDirichletBCValue(self, t, coords, node_id, primary_vars):
-		x, y, z = coords
-		
-		# prescribe displacement u_y
-		value = ExternalDisplacement(x, t)
-				
-		return (True, value)
+    def getDirichletBCValue(self, t, coords, node_id, primary_vars):
+        x, y, z = coords
+
+        # prescribe displacement u_y
+        value = ExternalDisplacement(x, t)
+
+        return (True, value)
+
 
 class BCM_BottomDisplacement_Y(OpenGeoSys.BoundaryCondition):
-	
-	def getDirichletBCValue(self, t, coords, node_id, primary_vars):
-		# prescribe displacement u_y
-		value = 0
-				
-		return (True, value)
+    def getDirichletBCValue(self, t, coords, node_id, primary_vars):
+        # prescribe displacement u_y
+        value = 0
+
+        return (True, value)
+
 
 class BCM_LateralDisplacement_X(OpenGeoSys.BoundaryCondition):
-	
-	def getDirichletBCValue(self, t, coords, node_id, primary_vars):
-		# prescribe displacement u_x
-		value = 0
-		
-		return (True, value)
+    def getDirichletBCValue(self, t, coords, node_id, primary_vars):
+        # prescribe displacement u_x
+        value = 0
+
+        return (True, value)
 
 
 # instantiate the BC objects used by OpenGeoSys
@@ -122,4 +129,3 @@ Traction4Monitoring_py = BCM_MonitoringTraction()
 LateralDisplacement_py = BCM_LateralDisplacement_X()
 SurfaceDisplacement_py = BCM_SurfaceDisplacement_Y()
 BottomDisplacement_py = BCM_BottomDisplacement_Y()
-