diff --git a/pycue/opencue/config.py b/pycue/opencue/config.py new file mode 100644 index 000000000..7d1040a87 --- /dev/null +++ b/pycue/opencue/config.py @@ -0,0 +1,96 @@ +# Copyright Contributors to the OpenCue Project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""OpenCue configuration.""" + +import logging +import os +import platform + +import yaml + + +logger = logging.getLogger("opencue") + + +# Config file from which default settings are loaded. This file is distributed with the +# opencue Python library. +__DEFAULT_CONFIG_FILE = os.path.join(os.path.dirname(__file__), 'default.yaml') + +# Environment variables which can be used to define a custom config file. Any settings +# defined in this file will be used instead of the defaults. +__CONFIG_FILE_ENV_VARS = [ + # OPENCUE_CONFIG_FILE is the preferred setting to use. + 'OPENCUE_CONFIG_FILE', + # OPENCUE_CONF is deprecated, but kept for now for backwards compatibility. + 'OPENCUE_CONF', +] + + +def config_base_directory(): + """Returns the OpenCue config base directory. + + This platform-dependent directory, stored within your user profile, is used by + OpenCue components as the default location for various configuration files. Typically + if you store your config files in this location, there is no need to set environment + variables to indicate where your config files are located -- OpenCue should recognize + them automatically. + + NOTE: This work is ongoing. Over time more OpenCue components will start using this + base directory. See https://github.com/AcademySoftwareFoundation/OpenCue/issues/785. + + :rtype: str + :return: config file base directory + """ + if platform.system() == 'Windows': + return os.path.join(os.path.expandvars('%APPDATA%'), 'opencue') + return os.path.join(os.path.expanduser('~'), '.config', 'opencue') + + +def load_config_from_file(): + """Loads configuration settings from config file on the local system. + + Default settings are read from default.yaml which is distributed with the opencue library. + User-provided config is then read from disk, in order of preference: + - Path defined by the OPENCUE_CONFIG_FILE environment variable. + - Path defined by the OPENCUE_CONF environment variable. + - Path within the config base directory (i.e. ~/.config/opencue/opencue.yaml) + + :rtype: dict + :return: config settings + """ + with open(__DEFAULT_CONFIG_FILE) as file_object: + config = yaml.load(file_object, Loader=yaml.SafeLoader) + + user_config_file = None + + for config_file_env_var in __CONFIG_FILE_ENV_VARS: + logger.debug('Checking for opencue config file path in %s', config_file_env_var) + config_file_from_env = os.environ.get(config_file_env_var) + if config_file_from_env and os.path.exists(config_file_from_env): + user_config_file = config_file_from_env + break + + if not user_config_file: + config_from_user_profile = os.path.join(config_base_directory(), 'opencue.yaml') + logger.debug('Checking for opencue config at %s', config_from_user_profile) + if os.path.exists(config_from_user_profile): + user_config_file = config_from_user_profile + + if user_config_file: + logger.info('Loading opencue config from %s', user_config_file) + with open(user_config_file) as file_object: + config.update(yaml.load(file_object, Loader=yaml.SafeLoader)) + + return config diff --git a/pycue/opencue/cuebot.py b/pycue/opencue/cuebot.py index 0e143e119..cb36d9d64 100644 --- a/pycue/opencue/cuebot.py +++ b/pycue/opencue/cuebot.py @@ -26,10 +26,10 @@ import logging import os import platform -import yaml import grpc +import opencue.config from opencue.compiled_proto import comment_pb2 from opencue.compiled_proto import comment_pb2_grpc from opencue.compiled_proto import criterion_pb2 @@ -67,15 +67,6 @@ logger = logging.getLogger("opencue") -default_config = os.path.join(os.path.dirname(__file__), 'default.yaml') -with open(default_config) as file_object: - config = yaml.load(file_object, Loader=yaml.SafeLoader) - -# check for facility specific configurations. -fcnf = os.environ.get('OPENCUE_CONF', '') -if os.path.exists(fcnf): - with open(fcnf) as file_object: - config.update(yaml.load(file_object, Loader=yaml.SafeLoader)) DEFAULT_MAX_MESSAGE_BYTES = 1024 ** 2 * 10 DEFAULT_GRPC_PORT = 8443 @@ -96,7 +87,8 @@ class Cuebot(object): RpcChannel = None Hosts = [] Stubs = {} - Timeout = config.get('cuebot.timeout', 10000) + Config = opencue.config.load_config_from_file() + Timeout = Config.get('cuebot.timeout', 10000) PROTO_MAP = { 'action': filter_pb2, @@ -150,13 +142,20 @@ class Cuebot(object): } @staticmethod - def init(): + def init(config=None): """Main init method for setting up the Cuebot object. - Sets the communication channel and hosts.""" + Sets the communication channel and hosts. + + :type config: dict + :param config: config dictionary, this will override the config read from disk + """ + if config: + Cuebot.Config = config + Cuebot.Timeout = config.get('cuebot.timeout', Cuebot.Timeout) if os.getenv("CUEBOT_HOSTS"): Cuebot.setHosts(os.getenv("CUEBOT_HOSTS").split(",")) else: - facility_default = config.get("cuebot.facility_default") + facility_default = Cuebot.Config.get("cuebot.facility_default") Cuebot.setFacility(facility_default) if Cuebot.Hosts is None: raise CueException('Cuebot host not set. Please ensure CUEBOT_HOSTS is set ' + @@ -169,7 +168,7 @@ def setChannel(): # gRPC must specify a single host. Randomize host list to balance load across cuebots. hosts = list(Cuebot.Hosts) shuffle(hosts) - maxMessageBytes = config.get('cuebot.max_message_bytes', DEFAULT_MAX_MESSAGE_BYTES) + maxMessageBytes = Cuebot.Config.get('cuebot.max_message_bytes', DEFAULT_MAX_MESSAGE_BYTES) # create interceptors interceptors = ( @@ -186,7 +185,8 @@ def setChannel(): if ':' in host: connectStr = host else: - connectStr = '%s:%s' % (host, config.get('cuebot.grpc_port', DEFAULT_GRPC_PORT)) + connectStr = '%s:%s' % ( + host, Cuebot.Config.get('cuebot.grpc_port', DEFAULT_GRPC_PORT)) logger.debug('connecting to gRPC at %s', connectStr) # TODO(bcipriano) Configure gRPC TLS. (Issue #150) try: @@ -228,12 +228,12 @@ def setFacility(facility): :type facility: str :param facility: a facility named in the config file""" - if facility not in list(config.get("cuebot.facility").keys()): - default = config.get("cuebot.facility_default") + if facility not in list(Cuebot.Config.get("cuebot.facility").keys()): + default = Cuebot.Config.get("cuebot.facility_default") logger.warning("The facility '%s' does not exist, defaulting to %s", facility, default) facility = default logger.debug("setting facility to: %s", facility) - hosts = config.get("cuebot.facility")[facility] + hosts = Cuebot.Config.get("cuebot.facility")[facility] Cuebot.setHosts(hosts) @staticmethod @@ -290,7 +290,7 @@ def getStub(cls, name): @staticmethod def getConfig(): """Gets the Cuebot config object, originally read in from the config file on disk.""" - return config + return Cuebot.Config # Python 2/3 compatible implementation of ABC diff --git a/pycue/tests/config_test.py b/pycue/tests/config_test.py new file mode 100644 index 000000000..9045365eb --- /dev/null +++ b/pycue/tests/config_test.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python + +# Copyright Contributors to the OpenCue Project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for `opencue.config`.""" + +import os +import unittest + +import mock +import pyfakefs.fake_filesystem_unittest + +import opencue.config + + +EXPECTED_DEFAULT_CONFIG = { + 'logger.format': '%(levelname)-9s %(module)-10s %(message)s', + 'logger.level': 'WARNING', + 'cuebot.protocol': 'tcp', + 'cuebot.grpc_port': 8443, + 'cuebot.timeout': 10000, + 'cuebot.max_message_bytes': 104857600, + 'cuebot.exception_retries': 3, + 'cuebot.facility_default': 'local', + 'cuebot.facility': { + 'local': ['localhost:8443'], + 'dev': ['cuetest02-vm.example.com:8443'], + 'cloud': [ + 'cuebot1.example.com:8443', + 'cuebot2.example.com:8443', + 'cuebot3.example.com:8443' + ], + }, +} + +USER_CONFIG = """ +cuebot.facility_default: fake-facility-01 +cuebot.facility: + fake-facility-01: + - fake-cuebot-01:1234 + fake-facility-02: + - fake-cuebot-02:5678 + - fake-cuebot-03:9012 +""" + + +class ConfigTests(pyfakefs.fake_filesystem_unittest.TestCase): + def setUp(self): + self.setUpPyfakefs() + self.fs.add_real_file( + os.path.join(os.path.dirname(opencue.__file__), 'default.yaml'), read_only=True) + os.unsetenv('OPENCUE_CONFIG_FILE') + os.unsetenv('OPENCUE_CONF') + + @mock.patch('platform.system', new=mock.Mock(return_value='Linux')) + @mock.patch('os.path.expanduser', new=mock.Mock(return_value='/home/username')) + def test__should_return_config_dir_unix(self): + self.assertEqual('/home/username/.config/opencue', opencue.config.config_base_directory()) + + @mock.patch('platform.system', new=mock.Mock(return_value='Windows')) + @mock.patch( + 'os.path.expandvars', new=mock.Mock(return_value='C:/Users/username/AppData/Roaming')) + def test__should_return_config_dir_windows(self): + self.assertEqual( + 'C:/Users/username/AppData/Roaming/opencue', opencue.config.config_base_directory()) + + def test__should_load_default_config(self): + self.assertIsNone(os.environ.get('OPENCUE_CONFIG_FILE')) + self.assertIsNone(os.environ.get('OPENCUE_CONF')) + + config = opencue.config.load_config_from_file() + + self.assertEqual(EXPECTED_DEFAULT_CONFIG, config) + + def test__should_load_user_config(self): + config_file_path = '/path/to/config.yaml' + self.fs.create_file(config_file_path, contents=USER_CONFIG) + os.environ['OPENCUE_CONFIG_FILE'] = config_file_path + # Define some invalid config using the old setting name, this ensures the old env var + # will be ignored if the new one is set. + config_file_path_legacy = '/path/to/legacy/config.yaml' + self.fs.create_file(config_file_path_legacy, contents='invalid yaml') + os.environ['OPENCUE_CONF'] = config_file_path_legacy + + config = opencue.config.load_config_from_file() + + self.assertEqual('fake-facility-01', config['cuebot.facility_default']) + self.assertEqual(['fake-cuebot-01:1234'], config['cuebot.facility']['fake-facility-01']) + self.assertEqual( + ['fake-cuebot-02:5678', 'fake-cuebot-03:9012'], + config['cuebot.facility']['fake-facility-02']) + # Settings not defined in user config should still have default values. + self.assertEqual(10000, config['cuebot.timeout']) + self.assertEqual(3, config['cuebot.exception_retries']) + + def test__should_load_user_config_from_legacy_var(self): + config_file_path = '/path/to/config.yaml' + self.fs.create_file(config_file_path, contents=USER_CONFIG) + os.environ['OPENCUE_CONF'] = config_file_path + + config = opencue.config.load_config_from_file() + + self.assertEqual('fake-facility-01', config['cuebot.facility_default']) + self.assertEqual(['fake-cuebot-01:1234'], config['cuebot.facility']['fake-facility-01']) + self.assertEqual( + ['fake-cuebot-02:5678', 'fake-cuebot-03:9012'], + config['cuebot.facility']['fake-facility-02']) + # Settings not defined in user config should still have default values. + self.assertEqual(10000, config['cuebot.timeout']) + self.assertEqual(3, config['cuebot.exception_retries']) + + @mock.patch('platform.system', new=mock.Mock(return_value='Linux')) + @mock.patch('os.path.expanduser', new=mock.Mock(return_value='/home/username')) + def test__should_load_user_config_from_user_profile(self): + config_file_path = '/home/username/.config/opencue/opencue.yaml' + self.fs.create_file(config_file_path, contents=USER_CONFIG) + + config = opencue.config.load_config_from_file() + + self.assertEqual('fake-facility-01', config['cuebot.facility_default']) + self.assertEqual(['fake-cuebot-01:1234'], config['cuebot.facility']['fake-facility-01']) + self.assertEqual( + ['fake-cuebot-02:5678', 'fake-cuebot-03:9012'], + config['cuebot.facility']['fake-facility-02']) + # Settings not defined in user config should still have default values. + self.assertEqual(10000, config['cuebot.timeout']) + self.assertEqual(3, config['cuebot.exception_retries']) + + +if __name__ == '__main__': + unittest.main() diff --git a/pycue/tests/cuebot_test.py b/pycue/tests/cuebot_test.py new file mode 100644 index 000000000..78e20fb97 --- /dev/null +++ b/pycue/tests/cuebot_test.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python + +# Copyright Contributors to the OpenCue Project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for `opencue.cuebot`.""" + +import unittest + +import mock + +import opencue + + +TESTING_CONFIG = { + "cuebot.facility_default": "fake-facility-01", + "cuebot.facility": { + "fake-facility-01": [ + "fake-cuebot-01", + ], + "fake-facility-02": [ + "fake-cuebot-02", + "fake-cuebot-03", + ], + }, +} + + +class CuebotTests(unittest.TestCase): + def setUp(self): + self.cuebot = opencue.Cuebot() + # Mocking the cue service ensures the initial healthcheck request made to Cuebot + # will succeed. + self.cuebot.SERVICE_MAP['cue'] = mock.Mock() + + def test__should_set_hosts_and_channel(self): + healthcheck_mock = mock.Mock() + self.cuebot.SERVICE_MAP['cue'] = healthcheck_mock + + self.cuebot.init(config=TESTING_CONFIG) + + self.assertEqual(["fake-cuebot-01"], self.cuebot.Hosts) + self.assertIsNotNone(self.cuebot.RpcChannel) + healthcheck_mock.assert_called_with(self.cuebot.RpcChannel) + + def test__should_set_known_facility(self): + self.cuebot.init(config=TESTING_CONFIG) + + self.cuebot.setFacility('fake-facility-02') + + self.assertEqual(['fake-cuebot-02', 'fake-cuebot-03'], self.cuebot.Hosts) + + def test__should_ignore_unknown_facility(self): + self.cuebot.init(config=TESTING_CONFIG) + + self.cuebot.setFacility('unknown-facility') + + self.assertEqual(['fake-cuebot-01'], self.cuebot.Hosts) + + +if __name__ == '__main__': + unittest.main()