diff --git a/rapyuta_io/clients/device.py b/rapyuta_io/clients/device.py index df96af2b..348e78d0 100644 --- a/rapyuta_io/clients/device.py +++ b/rapyuta_io/clients/device.py @@ -253,7 +253,8 @@ class Device(PartialMixin, ObjDict): def __init__(self, name, runtime=None, runtime_docker=False, runtime_preinstalled=False, ros_distro=None, rosbag_mount_path=None, - ros_workspace=None, description=None, python_version=DevicePythonVersion.PYTHON2): + ros_workspace=None, description=None, python_version=DevicePythonVersion.PYTHON2, + config_variables=None, labels=None): self.validate(name, runtime, runtime_docker, runtime_preinstalled, ros_distro, rosbag_mount_path, ros_workspace, description, python_version) self.name = name @@ -272,6 +273,8 @@ def __init__(self, name, runtime=None, runtime_docker=False, runtime_preinstalle self._ros_workspace = ros_workspace self.description = description self.python_version = python_version + self.config_variables = config_variables or {} + self.labels = labels or {} @staticmethod def validate(name, runtime, runtime_docker, runtime_preinstalled, ros_distro, rosbag_mount_path, @@ -320,6 +323,8 @@ def _serialize(self): '_ros_workspace']: if getattr(self, field): device['config_variables'][field[1:]] = getattr(self, field) + device['config_variables'].update(self.config_variables) + device['labels'] = self.labels return device @classmethod diff --git a/tests/device_test.py b/tests/device_test.py index 9b6742c0..bdc96b97 100644 --- a/tests/device_test.py +++ b/tests/device_test.py @@ -566,7 +566,8 @@ def test_create_device_dockercompose_success(self, mock_request): 'runtime_docker': True, 'ros_distro': 'melodic', 'rosbag_mount_path': 'test/path' - } + }, + 'labels': {} } expected_create_device_url = 'https://gaapiserver.apps.okd4v2.prod.rapyuta.io/api/device-manager/v0/auth-keys/?download_type=script' expected_get_device_url = 'https://gaapiserver.apps.okd4v2.prod.rapyuta.io/api/device-manager/v0/devices/test-device-id' @@ -610,7 +611,8 @@ def test_create_device_preinstalled_success(self, mock_request): 'runtime_preinstalled': True, 'ros_distro': 'melodic', 'ros_workspace': 'test/path' - } + }, + 'labels': {} } expected_create_device_url = 'https://gaapiserver.apps.okd4v2.prod.rapyuta.io/api/device-manager/v0/auth-keys/?download_type=script' expected_get_device_url = 'https://gaapiserver.apps.okd4v2.prod.rapyuta.io/api/device-manager/v0/devices/test-device-id' @@ -654,13 +656,22 @@ def test_create_device_dockercompose_success(self, mock_request): 'runtime_docker': True, 'runtime_preinstalled': True, 'ros_distro': 'melodic', - 'rosbag_mount_path': 'test/path' + 'rosbag_mount_path': 'test/path', + 'custom-config-variable-1': 'value1', + 'custom-config-variable-2': 'value2' + }, + 'labels': { + 'custom-label-1': 'label1', + 'custom-label-2': 'label2', + 'custom-label-3': 'label3' } } expected_create_device_url = 'https://gaapiserver.apps.okd4v2.prod.rapyuta.io/api/device-manager/v0/auth-keys/?download_type=script' expected_get_device_url = 'https://gaapiserver.apps.okd4v2.prod.rapyuta.io/api/device-manager/v0/devices/test-device-id' device = Device(name='test-device', runtime_docker=True, runtime_preinstalled=True, - ros_distro=ROSDistro.MELODIC, rosbag_mount_path='test/path', description='test-description') + ros_distro=ROSDistro.MELODIC, rosbag_mount_path='test/path', description='test-description', + config_variables={'custom-config-variable-1': 'value1','custom-config-variable-2': 'value2'}, + labels={'custom-label-1': 'label1','custom-label-2': 'label2','custom-label-3': 'label3'}) create_device_response = Mock() create_device_response.text = CREATE_BOTH_RUNTIMES_DEVICE_SUCCESS create_device_response.status_code = requests.codes.OK @@ -683,11 +694,21 @@ def test_create_device_dockercompose_success(self, mock_request): expected_configs = { 'runtime': 'dockercompose', 'ros_distro': 'melodic', - 'rosbag_mount_path': 'test/path' + 'rosbag_mount_path': 'test/path', + 'custom-config-variable-1': 'value1', + 'custom-config-variable-2': 'value2', } for config in device.config_variables: if config.key in expected_configs: self.assertEqual(expected_configs[config.key], config.value) + expected_labels = { + 'custom-label-1': 'label1', + 'custom-label-2': 'label2', + 'custom-label-3': 'label3' + } + for label in device.labels: + if label.key in expected_labels: + self.assertEqual(expected_labels[label.key], label.value) @patch('requests.request') def test_onboard_script_dockercompose_success(self, mock_request): @@ -893,7 +914,8 @@ def test_create_device_python3_dockercompose_success(self, mock_request): 'runtime_docker': True, 'ros_distro': 'melodic', 'rosbag_mount_path': 'test/path' - } + }, + 'labels': {} } expected_create_device_url = 'https://gaapiserver.apps.okd4v2.prod.rapyuta.io/api/device-manager/v0/auth-keys/?download_type=script' expected_get_device_url = 'https://gaapiserver.apps.okd4v2.prod.rapyuta.io/api/device-manager/v0/devices/test-device-id'