diff --git a/core/base/libs/socrate/socrate/system.py b/core/base/libs/socrate/socrate/system.py index b3b95012..99e11bbb 100644 --- a/core/base/libs/socrate/socrate/system.py +++ b/core/base/libs/socrate/socrate/system.py @@ -1,5 +1,6 @@ import socket import tenacity +from os import environ @tenacity.retry(stop=tenacity.stop_after_attempt(100), @@ -18,3 +19,14 @@ def resolve_address(address): hostname, *rest = address.rsplit(":", 1) ip_address = resolve_hostname(hostname) return ip_address + "".join(":" + port for port in rest) + + +def get_host_address_from_environment(name, default): + """ This function looks up an envionment variable ``{{ name }}_ADDRESS``. + If it's defined, it is returned unmodified. If it's undefined, an environment + variable ``HOST_{{ name }}`` is looked up and resolved to an ip address. + If this is also not defined, the default is resolved to an ip address. + """ + if "{}_ADDRESS".format(name) in environ: + return environ.get("{}_ADDRESS".format(name)) + return resolve_address(environ.get("HOST_{}".format(name), default)) diff --git a/core/base/libs/socrate/test.py b/core/base/libs/socrate/test.py index b9b36d5f..6fc87bfa 100644 --- a/core/base/libs/socrate/test.py +++ b/core/base/libs/socrate/test.py @@ -1,5 +1,6 @@ import unittest import io +import os from socrate import conf, system @@ -78,3 +79,30 @@ class TestSystem(unittest.TestCase): system.resolve_address("1.2.3.4.xip.io:80"), "1.2.3.4:80" ) + + def test_get_host_address_from_environment(self): + if "TEST_ADDRESS" in os.environ: + del os.environ["TEST_ADDRESS"] + if "HOST_TEST" in os.environ: + del os.environ["HOST_TEST"] + # if nothing is set, the default must be resolved + self.assertEqual( + system.get_host_address_from_environment("TEST", "1.2.3.4.xip.io:80"), + "1.2.3.4:80" + ) + # if HOST is set, the HOST must be resolved + os.environ['HOST_TEST']="1.2.3.5.xip.io:80" + self.assertEqual( + system.get_host_address_from_environment("TEST", "1.2.3.4.xip.io:80"), + "1.2.3.5:80" + ) + # if ADDRESS is set, the ADDRESS must be returned unresolved + os.environ['TEST_ADDRESS']="1.2.3.6.xip.io:80" + self.assertEqual( + system.get_host_address_from_environment("TEST", "1.2.3.4.xip.io:80"), + "1.2.3.6.xip.io:80" + ) + + +if __name__ == "__main__": + unittest.main()