def test_context_manager(self):
with for_each_client.for_each_client_backend(self._backend):
self.assertIs(for_each_client.get_for_each_client_backend(),
self._backend)
with for_each_client.for_each_client_backend('debug'):
self.assertIsInstance(
for_each_client.get_for_each_client_backend(),
for_each_client.ForEachClientDebugBackend)
self.assertIs(for_each_client.get_for_each_client_backend(),
self._backend)
def test_set_and_get_str(self):
with self.subTest('debug'):
for_each_client.set_for_each_client_backend('debug')
self.assertIsInstance(
for_each_client.get_for_each_client_backend(),
for_each_client.ForEachClientDebugBackend)
with self.subTest('jit'):
for_each_client.set_for_each_client_backend('jit')
self.assertIsInstance(
for_each_client.get_for_each_client_backend(),
for_each_client.ForEachClientJitBackend)
with self.subTest('pmap'):
for_each_client.set_for_each_client_backend('pmap')
self.assertIsInstance(
for_each_client.get_for_each_client_backend(),
for_each_client.ForEachClientPmapBackend)
with self.subTest('invalid'):
with self.assertRaisesRegex(ValueError,
"Unsupported backend 'invalid'"):
for_each_client.set_for_each_client_backend('invalid')
def test_set_and_get_concrete(self):
self.assertIsNot(for_each_client.get_for_each_client_backend(),
self._backend)
for_each_client.set_for_each_client_backend(self._backend)
self.assertIs(for_each_client.get_for_each_client_backend(),
self._backend)
def test_default_backend(self):
self.assertIsInstance(for_each_client.get_for_each_client_backend(),
for_each_client.ForEachClientJitBackend)