def test_incomplete_tf_config(self):
"""Test behavior when set TF_CONFIG is incomplete."""
with mock.patch.dict(os.environ, {'TF_CONFIG': '{"cluster": {}}'}):
cluster_config, save_dir = task.get_worker_behavior_info(
'some/dir/')
self.assertEqual(save_dir, 'some/dir/')
self.assertEqual(cluster_config, '')
with mock.patch.dict(os.environ, {'TF_CONFIG': '{"task": {}}'}):
cluster_config, save_dir = task.get_worker_behavior_info(
'some/dir/')
self.assertEqual(save_dir, 'some/dir/')
self.assertEqual(cluster_config, '')
def test_multi_worker_as_worker(self):
"""Tests multi-worker behavior when task type worker is set in TF_CONFIG."""
cluster_config, save_dir = task.get_worker_behavior_info('some/dir/')
self.assertEqual(save_dir, '')
self.assertEqual(
cluster_config,
('{"cluster": {"worker": ["worker0.example.com:2221"],'
'"chief": ["chief.example.com:2222"]},'
'"task": {"type": "worker", "index": 0}}'))
def test_single_worker(self):
"""Tests behavior when cluster has only one worker."""
cluster_config, save_dir = task.get_worker_behavior_info('some/dir/')
self.assertEqual(save_dir, 'some/dir/')
self.assertEqual(cluster_config, '')
def test_no_tf_config(self):
"""Tests behavior when there is no TF_CONFIG set."""
cluster_config, save_dir = task.get_worker_behavior_info('some/dir/')
self.assertEqual(save_dir, 'some/dir/')
self.assertEqual(cluster_config, '')