diff --git a/pyinfra/k8s_probes/startup.py b/pyinfra/k8s_probes/startup.py
index 72b8c54..986bbe8 100644
--- a/pyinfra/k8s_probes/startup.py
+++ b/pyinfra/k8s_probes/startup.py
@@ -29,7 +29,7 @@ def check_token_file():
         return False
 
 
-if __name__ == "__main__":
+def run_checks():
     if check_token_file():
         sys.exit(0)
     else: