async def test_writes_value(self, value, type_signature, expected_value):
root_dir = self.create_tempdir()
release_mngr = file_release_manager.SavedModelFileReleaseManager(
root_dir=root_dir, prefix='a_')
with mock.patch.object(file_utils,
'write_saved_model') as mock_write_saved_model:
await release_mngr.release(value, type_signature, 1)
mock_write_saved_model.assert_called_once()
call = mock_write_saved_model.mock_calls[0]
_, args, kwargs = call
actual_value, actual_path = args
def _normalize(value: Any) -> Any:
if isinstance(value, tf.data.Dataset):
return list(value)
return value
actual_value = tree.map_structure(_normalize, actual_value)
expected_value = tree.map_structure(_normalize, expected_value)
self.assertAllEqual(actual_value, expected_value)
expected_path = os.path.join(root_dir, 'a_1')
self.assertEqual(actual_path, expected_path)
self.assertEqual(kwargs, {'overwrite': True})
def test_raises_type_error_with_key(self, key):
root_dir = self.create_tempdir()
release_mngr = file_release_manager.SavedModelFileReleaseManager(
root_dir=root_dir, prefix='a_')
with self.assertRaises(TypeError):
release_mngr._get_path_for_key(key)
def test_returns_path_with_root_dir_and_prefix(self, root_dir, prefix, key,
expected_path):
release_mngr = file_release_manager.SavedModelFileReleaseManager(
root_dir=root_dir, prefix=prefix)
actual_path = release_mngr._get_path_for_key(key)
self.assertEqual(actual_path, expected_path)
def test_creates_new_dir_with_root_dir_path_like(self):
root_dir = self.create_tempdir()
shutil.rmtree(root_dir)
self.assertFalse(os.path.exists(root_dir))
file_release_manager.SavedModelFileReleaseManager(root_dir=root_dir)
self.assertTrue(os.path.exists(root_dir))
async def test_raises_type_error_with_key(self, key):
root_dir = self.create_tempdir()
release_mngr = file_release_manager.SavedModelFileReleaseManager(
root_dir=root_dir, prefix='a_')
value = 1
type_signature = computation_types.TensorType(tf.int32)
with self.assertRaises(TypeError):
await release_mngr.release(value, type_signature, key)
async def test_does_not_raise_type_error_with_key(self, key):
root_dir = self.create_tempdir()
release_mngr = file_release_manager.SavedModelFileReleaseManager(
root_dir=root_dir, prefix='a_')
value = 1
type_signature = computation_types.TensorType(tf.int32)
try:
await release_mngr.release(value, type_signature, key)
except TypeError:
self.fail('Raised TypeError unexpectedly.')
def test_raises_type_error_with_prefix(self, prefix):
root_dir = self.create_tempdir()
with self.assertRaises(TypeError):
file_release_manager.SavedModelFileReleaseManager(
root_dir=root_dir, prefix=prefix)
def test_raises_value_error_with_root_dir_empty(self):
with self.assertRaises(ValueError):
file_release_manager.SavedModelFileReleaseManager(root_dir='')
def test_raises_type_error_with_root_dir(self, root_dir):
with self.assertRaises(TypeError):
file_release_manager.SavedModelFileReleaseManager(
root_dir=root_dir)