diff --git a/yolo.py b/yolo.py
index 782d0c2..c1ad174 100755
--- a/yolo.py
+++ b/yolo.py
@@ -83,7 +83,7 @@ def get_os_friendly_name():
     else:
         return os_name
 
-def set_api_key():
+def set_api_key(config):
     """
     Set the OpenAI API key.
     """
@@ -107,35 +107,7 @@ def set_api_key():
     if not openai.api_key:
         openai.api_key = config["openai_api_key"]
 
-if __name__ == "__main__":
-    config = read_config()
-    set_api_key()
-
-    # Unix based SHELL (/bin/bash, /bin/zsh), otherwise assuming it's Windows
-    shell = os.environ.get("SHELL", "powershell.exe")
-
-    command_start_idx  = 1     # Question starts at which argv index?
-    ask_flag = False           # safety switch -a command line argument
-    yolo = ""                  # user's answer to safety switch (-a) question y/n
-
-    # Parse arguments and make sure we have at least a single word
-    if len(sys.argv) < 2:
-        print_usage()
-        sys.exit(-1)
-
-    # Safety switch via argument -a (local override of global setting)
-    # Force Y/n questions before running the command
-    if sys.argv[1] == "-a":
-        ask_flag = True
-        command_start_idx = 2
-
-    # To allow easy/natural use we don't require the input to be a
-    # single string. So, the user can just type yolo what is my name?
-    # without having to put the question between ''
-    arguments = sys.argv[command_start_idx:]
-    user_prompt = " ".join(arguments)
-
-def call_open_ai(query):
+def call_open_ai(config, shell, query):
     """
     Do we have a prompt from the user?
     """
@@ -187,7 +159,7 @@ def missing_posix_display():
     display = subprocess.check_output("echo $DISPLAY", shell=True)
     return display == b'\n'
 
-def prompt_user_input(response):
+def prompt_user_input(config, response):
     """
     Todo
     """
@@ -237,14 +209,40 @@ def main():
     """
     Defined starting point of source code.
     """
+    config = read_config()
+    set_api_key(config)
+
+    # Unix based SHELL (/bin/bash, /bin/zsh), otherwise assuming it's Windows
+    shell = os.environ.get("SHELL", "powershell.exe")
+
+    command_start_idx  = 1     # Question starts at which argv index?
+    ask_flag = False           # safety switch -a command line argument
+    yolo = ""                  # user's answer to safety switch (-a) question y/n
+
+    # Parse arguments and make sure we have at least a single word
+    if len(sys.argv) < 2:
+        print_usage()
+        sys.exit(-1)
+
+    # Safety switch via argument -a (local override of global setting)
+    # Force Y/n questions before running the command
+    if sys.argv[1] == "-a":
+        ask_flag = True
+        command_start_idx = 2
+
+    # To allow easy/natural use we don't require the input to be a
+    # single string. So, the user can just type yolo what is my name?
+    # without having to put the question between ''
+    arguments = sys.argv[command_start_idx:]
+    user_prompt = " ".join(arguments)
     
     #Enable color output on Windows using colorama
     init()
     
-    res_command = call_open_ai(user_prompt)
+    res_command = call_open_ai(config, shell, user_prompt)
     check_for_issue(res_command)
     check_for_markdown(res_command)
-    user_input = prompt_user_input(res_command)
+    user_input = prompt_user_input(config, res_command)
     print()
     evaluate_input(user_input, res_command)