diff --git a/judoscale/core/config.py b/judoscale/core/config.py index cc120f4..40913c3 100644 --- a/judoscale/core/config.py +++ b/judoscale/core/config.py @@ -45,7 +45,9 @@ def __init__( @classmethod def initialize(cls, env: Mapping = os.environ): - if env.get("DYNO"): + if env.get("JUDOSCALE_CONTAINER"): + return cls.for_custom(env) + elif env.get("DYNO"): return cls.for_heroku(env) elif env.get("RENDER_INSTANCE_ID"): return cls.for_render(env) @@ -56,7 +58,7 @@ def initialize(cls, env: Mapping = os.environ): elif env.get("RAILWAY_REPLICA_ID"): return cls.for_railway(env) else: - return cls.for_custom(env) + return cls.for_unknown(env) @classmethod def for_heroku(cls, env: Mapping): @@ -93,6 +95,12 @@ def for_railway(cls, env: Mapping): @classmethod def for_custom(cls, env: Mapping): + runtime_container = RuntimeContainer(env["JUDOSCALE_CONTAINER"]) + api_base_url = env.get("JUDOSCALE_URL") + return cls(runtime_container, api_base_url, env) + + @classmethod + def for_unknown(cls, env: Mapping): return cls(RuntimeContainer(""), env.get("JUDOSCALE_URL"), env) @property diff --git a/tests/test_config.py b/tests/test_config.py index 3bcc1f7..da6b5a1 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -83,6 +83,27 @@ def test_on_railway(self): assert config["API_BASE_URL"] == "https://adapter.judoscale.com/api/1234567890" def test_on_custom(self): + fake_env = { + "JUDOSCALE_CONTAINER": "my-custom-container", + "JUDOSCALE_URL": "https://adapter.judoscale.com/api/1234567890", + } + config = Config.initialize(fake_env) + + assert config["RUNTIME_CONTAINER"] == "my-custom-container" + assert config["LOG_LEVEL"] == "WARN" + assert config["API_BASE_URL"] == "https://adapter.judoscale.com/api/1234567890" + + def test_on_custom_takes_precedence(self): + fake_env = { + "JUDOSCALE_CONTAINER": "my-custom-container", + "DYNO": "web.1", + "JUDOSCALE_URL": "https://adapter.judoscale.com/api/1234567890", + } + config = Config.initialize(fake_env) + + assert config["RUNTIME_CONTAINER"] == "my-custom-container" + + def test_on_unknown(self): fake_env = {"JUDOSCALE_URL": "https://adapter.judoscale.com/api/1234567890"} config = Config.initialize(fake_env)