def testNeitherHostnameOrId(self):
with mock.patch.object(grr_colab.Client, 'with_id') as with_id_fn:
with mock.patch.object(grr_colab.Client,
'with_hostname') as with_hostname_fn:
with mock.patch.object(magics_impl, '_state', magics_impl._State()):
with self.assertRaises(ValueError):
magics_impl.grr_set_client_impl()
self.assertFalse(with_id_fn.called)
self.assertFalse(with_hostname_fn.called)
def testWithId(self):
mock_client = _MockClient(client_id='foo')
with mock.patch.object(
grr_colab.Client, 'with_id', return_value=mock_client) as with_id_fn:
with mock.patch.object(grr_colab.Client,
'with_hostname') as with_hostname_fn:
with mock.patch.object(magics_impl, '_state', magics_impl._State()):
magics_impl.grr_set_client_impl(client='foo')
self.assertFalse(with_hostname_fn.called)
with_id_fn.assert_called_once_with('foo')
self.assertEqual(magics_impl._state.client.id, mock_client.id)
