Merge pull request #10 from kilobyteno/check-required-env-vars

Author: dsbilling

tests/test_konfig.py

@@ -1,8 +1,9 @@
 import logging
+import os
 
 import pytest
 
-from tunsberg.konfig import uvicorn_log_config
+from tunsberg.konfig import check_required_env_vars, uvicorn_log_config
 
 
 class TestUvicornLogConfig:
@@ -55,3 +56,38 @@ def test_custom_log_level(self):
         assert config['loggers']['uvicorn']['level'] == logging.getLevelName(log_lvl)
         assert config['loggers']['uvicorn.error']['level'] == logging.getLevelName(log_lvl)
         assert config['loggers']['uvicorn.access']['level'] == logging.getLevelName(log_lvl)
+
+
+class TestCheckRequiredEnvVars:
+    # Set the environment variable for testing
+    os.environ['RANDOM_ENV_VAR'] = 'random_value'
+
+    def test_validate_if_true_in_local_development_env(self):
+        req_envs = {'RANDOM_ENV_VAR': {'runtime': True, 'build': True}}
+        check = check_required_env_vars(required_env_vars=req_envs, env='local')
+        assert check
+
+    def test_validate_if_true_in_staging_env(self):
+        req_envs = {'RANDOM_ENV_VAR': {'runtime': True, 'build': True}}
+        check = check_required_env_vars(required_env_vars=req_envs, env='staging', live_envs=['staging'])
+        assert check
+
+    def test_validate_if_true_in_prod_env(self):
+        req_envs = {'RANDOM_ENV_VAR': {'runtime': True, 'build': True}}
+        check = check_required_env_vars(required_env_vars=req_envs, env='prod', live_envs=['prod'])
+        assert check
+
+    def test_validate_that_it_fails_if_env_var_is_not_set(self):
+        req_envs = {'SOME_OTHER_RANDOM_ENV_VAR': {'runtime': True, 'build': True}}
+        with pytest.raises(ValueError):
+            check_required_env_vars(required_env_vars=req_envs, env='production', live_envs=['production'])
+
+    def test_validate_that_it_fails_for_code_build_if_env_var_is_not_set(self):
+        req_envs = {'SOME_OTHER_RANDOM_ENV_VAR': {'runtime': False, 'build': True}}
+        with pytest.raises(ValueError):
+            check_required_env_vars(required_env_vars=req_envs, env='production', live_envs=['production'], code_build=True)
+
+    def test_validate_that_it_fails_for_runtime_if_env_var_is_not_set(self):
+        req_envs = {'SOME_OTHER_RANDOM_ENV_VAR': {'runtime': True, 'build': False}}
+        with pytest.raises(ValueError):
+            check_required_env_vars(required_env_vars=req_envs, env='production', live_envs=['production'])

tunsberg/konfig.py

@@ -1,4 +1,5 @@
 import logging
+from os import getenv
 
 
 def uvicorn_log_config(
@@ -65,3 +66,49 @@ def uvicorn_log_config(
             },
         },
     }
+
+
+def check_required_env_vars(required_env_vars: dict, env: str, live_envs: [] or None = None, code_build: bool = False) -> bool or None:
+    """
+    Check if all required environment variables are set based on the current environment and if the code is being built.
+
+    Example for required_env_vars:
+        {'ENV': {'runtime': True, 'build': True},'JWT_PUBLIC_KEY': {'runtime': True, 'build': False}}
+
+    Example for env:
+        'production'
+
+    Example for live_envs:
+        ['production', 'prod', 'staging']
+
+    :param required_env_vars: Required environment variables
+    :type required_env_vars: dict
+    :param env: Current environment
+    :type env: str
+    :param live_envs: List of live environments
+    :type live_envs: [] or None
+    :param code_build: Whether the code is being built
+    :type code_build: bool
+    :return: True if all required environment variables are set
+    :rtype: bool or None
+    :raises Exception: If any required environment variable is not set
+    """
+    if live_envs is None:
+        live_envs = ['production', 'prod']
+
+    is_runtime = env in live_envs and not code_build
+    is_code_build = code_build
+
+    # Select environment variables based on current environment (production or build)
+    req_env_vars = [
+        env_var for env_var, conditions in required_env_vars.items() if (conditions['runtime'] and is_runtime) or (conditions['build'] and is_code_build)
+    ]
+
+    # Validate if all required environment variables are set
+    missing_vars = [env_var for env_var in req_env_vars if getenv(env_var) is None]
+
+    if missing_vars:
+        missing_vars_str = ', '.join(missing_vars)
+        raise ValueError(f'Environment Config Error: The following variables are not set: {missing_vars_str}')
+
+    return True