test_jobs_utils.py

fromunittestimportmock

fromsky.exceptionsimportClusterDoesNotExist
fromsky.jobsimportutils


@mock.patch('sky.core.down')
@mock.patch('sky.usage.usage_lib.messages.usage.set_internal')
deftest_terminate_cluster_retry_on_value_error(mock_set_internal,
mock_sky_down) ->None:
# Set up mock to fail twice with ValueError, then succeed
mock_sky_down.side_effect= [
ValueError('Mock error 1'),
ValueError('Mock error 2'),
None,
 ]

# Call should succeed after retries
utils.terminate_cluster('test-cluster')

# Verify sky.down was called 3 times
assertmock_sky_down.call_count==3
mock_sky_down.assert_has_calls([
mock.call('test-cluster'),
mock.call('test-cluster'),
mock.call('test-cluster'),
 ])

# Verify usage.set_internal was called before each sky.down
assertmock_set_internal.call_count==3


@mock.patch('sky.core.down')
@mock.patch('sky.usage.usage_lib.messages.usage.set_internal')
deftest_terminate_cluster_handles_nonexistent_cluster(mock_set_internal,
mock_sky_down) ->None:
# Set up mock to raise ClusterDoesNotExist
mock_sky_down.side_effect=ClusterDoesNotExist('test-cluster')

# Call should succeed silently
utils.terminate_cluster('test-cluster')

# Verify sky.down was called once
assertmock_sky_down.call_count==1
mock_sky_down.assert_called_once_with('test-cluster')

# Verify usage.set_internal was called once
assertmock_set_internal.call_count==1