diff --git a/tests/test_devices.py b/tests/test_devices.py index 7a01d8f7e..5d792826c 100644 --- a/tests/test_devices.py +++ b/tests/test_devices.py @@ -3,10 +3,18 @@ from visual_behavior import devices -@pytest.mark.parametrize("kwargs, expected", [ - ({"in_val": "W7DTMJ19R2F", "input_type": "computer_name", }, "A1", ), - ({"in_val": "A2", "input_type": "rig_id", }, "W7DTMJ35Y0T", ), - ({"in_val": "NOT_A_NAME", "input_type": "computer_name", }, "unknown", ), +@pytest.mark.parametrize("computer_name, rig_id", [ + ("W7DTMJ19R2F", "A1"), + ("w7dtmj19r2f", "A1"), + ("NOT_A_NAME", "unknown"), ]) -def test_get_rid_id(kwargs, expected): - assert devices.get_rig_id(**kwargs) == expected +def test_get_rid_id(computer_name, rig_id): + assert devices.get_rig_id(computer_name) == rig_id + + +@pytest.mark.parametrize("rig_id, computer_name", [ + ("A2","w7dtmj35y0t"), + ("NOT_A_RIG", "unknown"), +]) +def test_get_computer_name(rig_id, computer_name): + assert devices.get_computer_name(rig_id) == computer_name diff --git a/visual_behavior/devices.py b/visual_behavior/devices.py index 7e887db05..87b9fec77 100644 --- a/visual_behavior/devices.py +++ b/visual_behavior/devices.py @@ -51,29 +51,39 @@ 'localhost': 'localhost' } +RIG_NAME = {k.lower(): v for k, v in iteritems(RIG_NAME)} + COMPUTER_NAME = dict((v, k) for k, v in iteritems(RIG_NAME)) # -> devices.py -def get_rig_id(in_val, input_type='computer_name'): +def get_rig_id(computer_name): ''' This provides a map between the computer name and the rig ID. >>> get_rig_id('W7DTMJ19R2F') A1 - >>> get_rig_id('A1',input_type='rig_id') - W7DTMJ19R2F Parameters ---------- in_val : str computer name - input_type : 'computer_name' or 'rig_id' - specifies whether `in_val` is a computer name or rig_id ''' - if input_type == 'computer_name' and in_val in RIG_NAME.keys(): - return RIG_NAME[in_val] - elif input_type == 'rig_id' and in_val in COMPUTER_NAME.keys(): - return COMPUTER_NAME[in_val] - else: - return 'unknown' + + return RIG_NAME.get(computer_name.lower(), 'unknown') + + +def get_computer_name(rig_id): + ''' + This provides a map between the computer name and the rig ID. + + >>> get_computer_name('A1') + W7DTMJ19R2F + + Parameters + ---------- + rig_id : str + rig name + ''' + + return COMPUTER_NAME.get(rig_id, 'unknown')