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)