def test_handle_multiple_roles(self, _config_mock):
keyman = Keyman(["foo", "-o", "foo", "-u", "bar", "-a", "baz"])
keyman.selector_menu = mock.MagicMock()
keyman.selector_menu.return_value = 0
roles = [
{"account": "acct1", "role_name": "role1", "roleIdx": 0},
{"account": "acct1", "role_name": "role2", "roleIdx": 1},
]
mock_session = mock.MagicMock()
mock_session.available_roles.return_value = roles
keyman.config.account = None
keyman.config.role = None
keyman.handle_multiple_roles(mock_session)
keyman.selector_menu.assert_has_calls(
[
mock.call(
roles,
[{"account": "Account"}, {"role_name": "Role"}],
),
],
)
mock_session.assert_has_calls(
[
mock.call.available_roles(),
],
)
def test_handle_multiple_roles(self, _config_mock):
keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz'])
keyman.selector_menu = mock.MagicMock()
keyman.selector_menu.return_value = 0
roles = ([{
'account': 'acct1',
'role_name': 'role1',
'roleIdx': 0
}, {
'account': 'acct1',
'role_name': 'role2',
'roleIdx': 1
}])
mock_session = mock.MagicMock()
mock_session.available_roles.return_value = roles
keyman.config.account = None
keyman.config.role = None
keyman.handle_multiple_roles(mock_session)
keyman.selector_menu.assert_has_calls([
mock.call(roles, [{
'account': 'Account'
}, {
'role_name': 'Role'
}])
])
mock_session.assert_has_calls([mock.call.available_roles()])
def test_handle_multiple_roles_single_role_match(self, _config_mock):
keyman = Keyman(["foo", "-o", "foo", "-u", "bar", "-a", "baz"])
keyman.selector_menu = mock.MagicMock()
keyman.selector_menu.return_value = 0
roles = [
{"account": "acct1", "role_name": "role1", "roleIdx": 0},
{"account": "acct1", "role_name": "role2", "roleIdx": 1},
{"account": "acct2", "role_name": "role1", "roleIdx": 2},
{"account": "acct2", "role_name": "role2", "roleIdx": 3},
]
mock_session = mock.MagicMock()
mock_session.available_roles.return_value = roles
keyman.config.account = "acct1"
keyman.config.role = "role2"
assert keyman.handle_multiple_roles(mock_session)
assert keyman.role == 1
keyman.selector_menu.assert_not_called()
mock_session.assert_has_calls(
[
mock.call.available_roles(),
],
)
def test_handle_multiple_roles_no_match_role(self, _config_mock):
keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz'])
keyman.selector_menu = mock.MagicMock()
keyman.selector_menu.return_value = 0
roles = ([{
'account': 'acct1',
'role_name': 'role1',
'roleIdx': 0
}, {
'account': 'acct1',
'role_name': 'role2',
'roleIdx': 1
}, {
'account': 'acct2',
'role_name': 'role1',
'roleIdx': 2
}])
mock_session = mock.MagicMock()
mock_session.available_roles.return_value = roles
keyman.config.account = 'role3'
keyman.config.role = None
assert not keyman.handle_multiple_roles(mock_session)
keyman.selector_menu.assert_not_called()
mock_session.assert_has_calls([mock.call.available_roles()])
def test_aws_auth_loop_multirole(self, config_mock):
config_mock().reup = False
keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz'])
keyman.start_session = mock.MagicMock()
keyman.start_session().assume_role.side_effect = aws.MultipleRoles
keyman.handle_multiple_roles = mock.MagicMock()
keyman.aws_auth_loop()
keyman.handle_multiple_roles.assert_has_calls([mock.call(mock.ANY)])
def test_aws_auth_loop(self, config_mock):
config_mock().reup = False
keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz'])
keyman.start_session = mock.MagicMock()
keyman.wrap_up = mock.MagicMock()
keyman.handle_multiple_roles = mock.MagicMock()
keyman.handle_multiple_roles.return_value = True
keyman.aws_auth_loop()
keyman.start_session.assert_has_calls(
[mock.call(), mock.call().assume_role(mock.ANY)])
assert keyman.wrap_up.called
assert keyman.handle_multiple_roles.called
def test_aws_auth_loop_reauth(self, config_mock, _sleep_mock):
config_mock().reup = False
keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz'])
session_instance = mock.MagicMock()
session_instance.assume_role.side_effect = okta.ReauthNeeded
session_instance.is_valid = False
keyman.start_session = mock.MagicMock()
keyman.start_session.side_effect = session_instance, Exception()
keyman.handle_multiple_roles = mock.MagicMock()
keyman.auth_okta = mock.MagicMock()
with self.assertRaises(Exception):
keyman.aws_auth_loop()
keyman.auth_okta.assert_has_calls([mock.call(state_token=None)])
def test_handle_multiple_roles_rollname_match(self, _config_mock):
keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz'])
keyman.selector_menu = mock.MagicMock()
keyman.selector_menu.return_value = 0
roles = ([{
'account': 'acct1',
'role_name': 'role1',
'roleIdx': 0
}, {
'account': 'acct1',
'role_name': 'role2',
'roleIdx': 1
}, {
'account': 'acct2',
'role_name': 'role1',
'roleIdx': 2
}, {
'account': 'acct2',
'role_name': 'role2',
'roleIdx': 3
}])
mock_session = mock.MagicMock()
mock_session.available_roles.return_value = roles
keyman.config.account = None
keyman.config.role = 'role2'
assert keyman.handle_multiple_roles(mock_session)
# item 0 "selected" from menu translates to role #1
assert keyman.role == 1
keyman.selector_menu.assert_has_calls([
mock.call(([{
'account': 'acct1',
'role_name': 'role2',
'roleIdx': 1
}, {
'account': 'acct2',
'role_name': 'role2',
'roleIdx': 3
}]), [{
'account': 'Account'
}, {
'role_name': 'Role'
}])
])
mock_session.assert_has_calls([mock.call.available_roles()])
def test_aws_auth_loop_multirole_no_match(self, config_mock):
config_mock().reup = False
config_mock().screen = False
keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz'])
session_instance = mock.MagicMock()
session_instance.assume_role.side_effect = aws.MultipleRoles
keyman.start_session = mock.MagicMock()
keyman.start_session.return_value = session_instance
keyman.handle_multiple_roles = mock.MagicMock()
keyman.handle_multiple_roles.return_value = False
ret = keyman.aws_auth_loop()
keyman.handle_multiple_roles.assert_has_calls([mock.call(mock.ANY)])
self.assertEqual(ret, 1)
def test_aws_auth_loop_continue(self, config_mock, sleep_mock):
config_mock().reup = False
keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz'])
keyman.config.reup = True
session_instance = mock.MagicMock()
session_instance.is_valid = True
keyman.start_session = mock.MagicMock()
keyman.start_session.return_value = session_instance
keyman.handle_multiple_roles = mock.MagicMock()
keyman.handle_multiple_roles.return_value = True
with self.assertRaises(Exception):
keyman.aws_auth_loop()
sleep_mock.assert_has_calls([mock.call(60)])
assert keyman.handle_multiple_roles.called
def test_aws_auth_loop(self, config_mock):
config_mock().reup = False
keyman = Keyman(["foo", "-o", "foo", "-u", "bar", "-a", "baz"])
keyman.start_session = mock.MagicMock()
keyman.wrap_up = mock.MagicMock()
keyman.handle_multiple_roles = mock.MagicMock()
keyman.handle_multiple_roles.return_value = True
keyman.aws_auth_loop()
keyman.start_session.assert_has_calls(
[
mock.call(),
mock.call().assume_role(mock.ANY),
],
)
assert keyman.wrap_up.called
assert keyman.handle_multiple_roles.called
def test_aws_auth_loop_multirole(self, config_mock):
config_mock().reup = False
keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz'])
session_instance = mock.MagicMock()
session_instance.assume_role.side_effect = aws.MultipleRoles
keyman.start_session = mock.MagicMock()
keyman.start_session.return_value = session_instance
keyman.handle_multiple_roles = mock.MagicMock()
with self.assertRaises(Exception):
keyman.aws_auth_loop()
session_instance.assume_role.assert_has_calls([
mock.call(),
mock.call(),
])
keyman.handle_multiple_roles.assert_has_calls([
mock.call(mock.ANY)
])
def test_handle_multiple_roles(self, _config_mock):
keyman = Keyman(['foo', '-o', 'foo', '-u', 'bar', '-a', 'baz'])
keyman.selector_menu = mock.MagicMock()
keyman.selector_menu.return_value = 0
roles = [{}, {}]
mock_session = mock.MagicMock()
mock_session.available_roles.return_value = roles
ret = keyman.handle_multiple_roles(mock_session)
self.assertEqual(ret, 0)
keyman.selector_menu.assert_has_calls(
[mock.call([{}, {}], 'role', 'Role')])
mock_session.assert_has_calls([
mock.call.available_roles(),
mock.call.set_role(mock.ANY),
mock.call.assume_role()
])
def test_handle_multiple_roles_rollname_match(self, _config_mock):
keyman = Keyman(["foo", "-o", "foo", "-u", "bar", "-a", "baz"])
keyman.selector_menu = mock.MagicMock()
keyman.selector_menu.return_value = 0
roles = [
{"account": "acct1", "role_name": "role1", "roleIdx": 0},
{"account": "acct1", "role_name": "role2", "roleIdx": 1},
{"account": "acct2", "role_name": "role1", "roleIdx": 2},
{"account": "acct2", "role_name": "role2", "roleIdx": 3},
]
mock_session = mock.MagicMock()
mock_session.available_roles.return_value = roles
keyman.config.account = None
keyman.config.role = "role2"
assert keyman.handle_multiple_roles(mock_session)
# item 0 "selected" from menu translates to role #1
assert keyman.role == 1
keyman.selector_menu.assert_has_calls(
[
mock.call(
(
[
{"account": "acct1", "role_name": "role2", "roleIdx": 1},
{"account": "acct2", "role_name": "role2", "roleIdx": 3},
]
),
[{"account": "Account"}, {"role_name": "Role"}],
),
],
)
mock_session.assert_has_calls(
[
mock.call.available_roles(),
],
)
