common_test_fixtures.py

importbase64
importcollections
importos
importpickle
importtempfile
importtime
importunittest
importuuid

importboto3
importfastapi
fromfastapiimporttestclient
importpandasaspd
importpytest
importrequests

importsky
fromskyimportglobal_user_state
fromskyimportsky_logging
fromsky.backends.cloud_vm_ray_backendimportCloudVmRayBackend
fromsky.clouds.service_catalogimportvsphere_catalog
fromsky.provisionimportcommonasprovision_common
fromsky.provision.awsimportconfigasaws_config
fromsky.provision.kubernetesimportutilsaskubernetes_utils
fromsky.serveimportserve_state
fromsky.serverimportcommonasserver_common
fromsky.server.requestsimportexecutor
fromsky.server.requestsimportrequestsasapi_requests
fromsky.server.serverimportapp
fromsky.skyletimportconstants
fromsky.utilsimportcontroller_utils
fromsky.utilsimportmessage_utils
fromsky.utilsimportregistry

logger=sky_logging.init_logger("sky.pytest")


@pytest.fixture
defaws_config_region(monkeypatch: pytest.MonkeyPatch) ->str:
fromskyimportskypilot_config
region='us-east-2'
ifskypilot_config.loaded():
ssh_proxy_command=skypilot_config.get_nested(
 ('aws', 'ssh_proxy_command'), None)
ifisinstance(ssh_proxy_command, dict) andssh_proxy_command:
region=list(ssh_proxy_command.keys())[0]
returnregion


@pytest.fixture
defmock_client_requests(monkeypatch: pytest.MonkeyPatch, mock_queue,
mock_stream_utils, mock_redirect_log_file) ->None:
"""Fixture to mock HTTP requests using FastAPI's TestClient."""
# This fixture automatically replaces `requests.get` and `requests.post`
# with mocked versions that route requests through a FastAPI TestClient.
# It is used to simulate server responses for testing purposes without
# making actual HTTP requests.
client=testclient.TestClient(app)
original_requests_get=requests.get
original_requests_post=requests.post

def_execute_request(path: str, method: str,
response: fastapi.Response) ->None:
request_id=server_common.get_request_id(response)
request_obj=api_requests.get_request(request_id)
# To have the mock effective, we do not start the backend request
# workers, i.e. the requests placed in the database will not be
# executed automatically. We manually call the executor._wrapper
# here to execute the request.
ifrequest_objisnotNone:
ignore_return_value=mock_queue.get(request_id)
ifignore_return_valueisNoneandpath=='/optimize':
ignore_return_value=True
logger.info(
f'Mocking {method} request to {path} with request_id {request_id}, '
f'running executor._wrapper(\'{request_id}\', {ignore_return_value})'
 )
executor._request_execution_wrapper(request_id, ignore_return_value)

defmock_http_request(method: str, url, *args, **kwargs):
ifmethod=='GET':
mock_func=client.get
original_func=original_requests_get
elifmethod=='POST':
mock_func=client.post
original_func=original_requests_post
else:
raiseValueError(f'Unsupported method: {method}')
ifserver_common.get_server_url() inurl:
logger.info(f'Mocking {method} request to {url} through TestClient')
path=url.replace(server_common.get_server_url(), "")
# Remove stream parameter as it's not supported by TestClient
stream=kwargs.pop('stream', False)
# Extract and format query parameters
if'params'inkwargs:
kwargs['params'] = {
k: vfork, vinkwargs['params'].items() ifvisnotNone
 }
response=mock_func(path, *args, **kwargs)
ifnotany(
path.startswith(p)
forpin ['/api/get', '/api/stream', '/api/status']):
# These paths do not need to be executed
_execute_request(path, method, response)
# If streaming is requested, wrap the response content in an iterator
ifstream:
content=response.content

defiter_content(chunk_size=None):
# Yield the entire content as one chunk since this is a test
ifnotcontent:
yieldNone
yieldcontent

# Add iter_content method to response
response.iter_content=iter_content

returnresponse
else:
returnoriginal_func(url, *args, **kwargs)

# Mock `requests.post` to use TestClient.post
defmock_post(url, *args, **kwargs):
# Convert full URL to path for TestClient
returnmock_http_request('POST', url, *args, **kwargs)

# Mock `requests.get` to use TestClient.get
defmock_get(url, *args, **kwargs):
returnmock_http_request('GET', url, *args, **kwargs)

# Use monkeypatch to replace `requests.post` and `requests.get`
monkeypatch.setattr(requests, "post", mock_post)
monkeypatch.setattr(requests, "get", mock_get)


@pytest.fixture
defenable_all_clouds(monkeypatch, request, mock_client_requests):
"""Create mock context managers for cloud configurations."""
enabled_clouds=request.paramifhasattr(request, 'param') elseNone
ifenabled_cloudsisNone:
enabled_clouds=list(registry.CLOUD_REGISTRY.values())

config_file=tempfile.NamedTemporaryFile(prefix='tmp_config_default',
delete=False).name

# Mock all the functions
monkeypatch.setattr('sky.check.get_cached_enabled_clouds_or_refresh',
lambda*_, **__: enabled_clouds)
monkeypatch.setattr('sky.check.check_capability', lambda*_, **__: None)
monkeypatch.setattr(
'sky.clouds.service_catalog.aws_catalog._get_az_mappings',
lambda*_, **__: pd.read_csv('tests/default_aws_az_mappings.csv'))
monkeypatch.setattr('sky.backends.backend_utils.check_owner_identity',
lambda*_, **__: None)
monkeypatch.setattr(
'sky.clouds.utils.gcp_utils.list_reservations_for_instance_type_in_zone',
lambda*_, **__: [])

# Kubernetes mocks
monkeypatch.setattr('sky.adaptors.kubernetes._load_config',
lambda*_, **__: None)
monkeypatch.setattr(
'sky.provision.kubernetes.utils.detect_gpu_label_formatter',
lambda*_, **__: [kubernetes_utils.SkyPilotLabelFormatter, {}])
monkeypatch.setattr(
'sky.provision.kubernetes.utils.detect_accelerator_resource',
lambda*_, **__: [True, []])
monkeypatch.setattr('sky.provision.kubernetes.utils.check_instance_fits',
lambda*_, **__: [True, ''])
monkeypatch.setattr('sky.provision.kubernetes.utils.get_spot_label',
lambda*_, **__: [None, None])
monkeypatch.setattr('sky.clouds.kubernetes.kubernetes_utils.get_spot_label',
lambda*_, **__: [None, None])
monkeypatch.setattr(
'sky.provision.kubernetes.utils.is_kubeconfig_exec_auth',
lambda*_, **__: [False, None])
monkeypatch.setattr(
'sky.clouds.kubernetes.Kubernetes.regions_with_offering',
lambda*_, **__: [sky.clouds.Region('my-k8s-cluster-context')])

# VSphere catalog mock
monkeypatch.setattr(vsphere_catalog, '_LOCAL_CATALOG',
'tests/default_vsphere_vms.csv')

# Mock quota checking for enabled clouds
forcloudinenabled_clouds:
ifhasattr(cloud, 'check_quota_available'):
monkeypatch.setattr(cloud, 'check_quota_available',
lambda*_, **__: True)

# Environment variables
monkeypatch.setattr(
'sky.clouds.gcp.DEFAULT_GCP_APPLICATION_CREDENTIAL_PATH', config_file)
monkeypatch.setenv('OCI_CONFIG', config_file)


@pytest.fixture
defmock_job_table_no_job(monkeypatch):
"""Mock job table to return no jobs."""

defmock_get_job_table(*_, **__):
return0, message_utils.encode_payload([]), ''

monkeypatch.setattr(CloudVmRayBackend, 'run_on_head', mock_get_job_table)


@pytest.fixture
defmock_job_table_one_job(monkeypatch):
"""Mock job table to return one job."""

defmock_get_job_table(*_, **__):
current_time=time.time()
job_data= {
'job_id': '1',
'job_name': 'test_job',
'resources': 'test',
'status': 'RUNNING',
'submitted_at': current_time,
'run_timestamp': str(current_time),
'start_at': current_time,
'end_at': current_time,
'last_recovered_at': None,
'recovery_count': 0,
'failure_reason': '',
'managed_job_id': '1',
'task_id': 0,
'task_name': 'test_task',
'job_duration': 20,
 }
return0, message_utils.encode_payload([job_data]), ''

monkeypatch.setattr(CloudVmRayBackend, 'run_on_head', mock_get_job_table)


@pytest.fixture
defmock_controller_accessible(monkeypatch):
"""Mock controller to be accessible."""

defmock_is_controller_accessible(controller: controller_utils.Controllers,
*_, **__):
record=global_user_state.get_cluster_from_name(
controller.value.cluster_name)
returnrecord['handle'] # type: ignore

monkeypatch.setattr('sky.backends.backend_utils.is_controller_accessible',
mock_is_controller_accessible)


@pytest.fixture
defmock_services_no_service(monkeypatch):
"""Mock services to return no services."""

defmock_get_services(*_, **__):
return0, message_utils.encode_payload(
 [], payload_type='service_status'), ''

monkeypatch.setattr(CloudVmRayBackend, 'run_on_head', mock_get_services)


@pytest.fixture
defmock_services_one_service(monkeypatch):
"""Mock services to return one service."""

defmock_get_services(*_, **__):
service= {
'name': 'test_service',
'controller_job_id': 1,
'uptime': 20,
'status': serve_state.ServiceStatus.READY,
'controller_port': 30001,
'load_balancer_port': 30000,
'endpoint': '4.3.2.1:30000',
'policy': None,
'requested_resources_str': '',
'replica_info': [],
'tls_encrypted': False,
 }
return0, message_utils.encode_payload(
 [{
k: base64.b64encode(pickle.dumps(v)).decode('utf-8')
fork, vinservice.items()
 }],
payload_type='service_status'), ''

monkeypatch.setattr(CloudVmRayBackend, 'run_on_head', mock_get_services)


@pytest.fixture
defmock_queue(monkeypatch):
"""
 Mock for `sky.server.requests.queues.mp_queue.get_queue` to return an object
 with a `.put` method that stores (request_id, ignore_return_value) in a map.
 """

# Define a mock queue object with a `.put` method
classMockQueue:

def__init__(self):
# Store (request_id, ignore_return_value) pairs in a dictionary
self.queue_map= {}

defput(self, item):
# Add to the map; item is assumed to be a tuple (request_id, ignore_return_value)
request_id, ignore_return_value=item
self.queue_map[request_id] =ignore_return_value

defget(self, request_id):
# Retrieve ignore_return_value for a given request_id
returnself.queue_map.get(request_id)

# Create a MockQueue instance
mock_queue_instance=MockQueue()

# Mock `get_queue` to return the mock_queue_instance
defmock_get_queue(schedule_type):
returnmock_queue_instance

# Apply monkeypatch to replace `mp_queue.get_queue`
monkeypatch.setattr("sky.server.requests.queues.mp_queue.get_queue",
mock_get_queue)

# Return the mock_queue_instance for use in tests
returnmock_queue_instance


@pytest.fixture
defmock_redirect_log_file(monkeypatch):
monkeypatch.setattr('sky.server.requests.executor._redirect_output',
lambda*_, **__: (None, None))
monkeypatch.setattr('sky.server.requests.executor._restore_output',
lambda*_, **__: None)


@pytest.fixture
defmock_stream_utils(monkeypatch):
# Patch out the original stream_utils.stream_response so it returns None
# rather than a StreamingResponse—this ensures the call won't block.
def_mock_empty_generator():
yieldb""

def_mock_stream_response(*args, **kwargs):
returnfastapi.responses.StreamingResponse(
_mock_empty_generator(),
media_type="text/plain",
headers={
'Cache-Control': 'no-cache, no-transform',
'X-Accel-Buffering': 'no',
'Transfer-Encoding': 'chunked'
 })

monkeypatch.setattr("sky.server.stream_utils.stream_response",
_mock_stream_response)


@pytest.fixture
defmock_aws_backend(monkeypatch):
"""Mock AWS backend for basic AWS testing operations."""
# Create a Subnet class to match what SkyPilot expects
Subnet=collections.namedtuple(
'Subnet',
 ['subnet_id', 'vpc_id', 'availability_zone', 'map_public_ip_on_launch'])

# Mock subnet and VPC discovery
defmock_get_subnet_and_vpc_id(*args, **kwargs):
# Get the region from kwargs
region=kwargs.get('region', 'us-east-1')

# Create a subnet in the requested region
ec2=boto3.resource('ec2', region_name=region)

# Create VPC and subnet in the requested region
vpc=ec2.create_vpc(CidrBlock='10.0.0.0/16')
subnet=ec2.create_subnet(VpcId=vpc.id,
CidrBlock='10.0.0.0/24',
AvailabilityZone=f"{region}a")

# Configure subnet to map public IPs
ec2.meta.client.modify_subnet_attribute(
SubnetId=subnet.id, MapPublicIpOnLaunch={'Value': True})

# Store subnet object
subnet_obj=Subnet(subnet_id=subnet.id,
vpc_id=vpc.id,
availability_zone=f"{region}a",
map_public_ip_on_launch=True)

return ([subnet_obj], vpc.id)

# Mock security groups
defmock_get_or_create_vpc_security_group(*args, **kwargs):
# Return a mock security group
returnunittest.mock.Mock(id="sg-12345678", group_name="test-sg")

# Mock IAM role
defmock_configure_iam_role(*args, **kwargs):
return {'Name': 'skypilot-test-role'}

defmock_wait_instances(region, cluster_name_on_cloud, state):
# Always return successfully without waiting
return

defmock_post_provision_runtime_setup(cloud_name, cluster_name,
cluster_yaml, provision_record,
custom_resource, log_dir):
# Get region from the provision record
region=provision_record.region

# Create a head instance
head_instance_id=f'i-{uuid.uuid4().hex[:8]}'

# Create instance info for the head
head_instance=provision_common.InstanceInfo(
instance_id=head_instance_id,
internal_ip='10.0.0.1',
external_ip='192.168.1.1',
tags={
'Name': cluster_name.name_on_cloud,
'ray-cluster-name': cluster_name.name_on_cloud,
'ray-node-type': 'head'
 })

# Create ClusterInfo
instances= {head_instance_id: [head_instance]}
cluster_info=provision_common.ClusterInfo(
instances=instances,
head_instance_id=head_instance_id,
provider_name='aws',
provider_config={
'region': region,
'use_internal_ips': False
 },
ssh_user='ubuntu')

returncluster_info

defmock_execute(self, handle, task, detach_run, dryrun=False):
# Return a fake job ID without attempting to SSH
return1234

# Apply our mocks to the monkeypatch
monkeypatch.setattr(aws_config, '_get_subnet_and_vpc_id',
mock_get_subnet_and_vpc_id)
monkeypatch.setattr(aws_config, '_get_or_create_vpc_security_group',
mock_get_or_create_vpc_security_group)
monkeypatch.setattr(aws_config, '_configure_iam_role',
mock_configure_iam_role)
monkeypatch.setattr(sky.provision.aws, 'wait_instances',
mock_wait_instances)
# Add mock for post_provision_runtime_setup
monkeypatch.setattr(sky.provision.provisioner,
'post_provision_runtime_setup',
mock_post_provision_runtime_setup)
# Add mock for _execute
monkeypatch.setattr(sky.backends.cloud_vm_ray_backend.CloudVmRayBackend,
'_execute', mock_execute)


@pytest.fixture
defskyignore_dir():
withtempfile.TemporaryDirectory() astemp_dir:
# Create workdir
dirs= ['remove_dir', 'dir', 'dir/subdir', 'dir/subdir/remove_dir']
files= [
'remove.py',
'remove.sh',
'remove.a',
'keep.py',
'remove.a',
'dir/keep.txt',
'dir/remove.sh',
'dir/keep.a',
'dir/remove.b',
'dir/remove.a',
'dir/subdir/keep.b',
'dir/subdir/remove.py',
 ]
fordir_nameindirs:
os.makedirs(os.path.join(temp_dir, dir_name), exist_ok=True)
forfile_pathinfiles:
full_path=os.path.join(temp_dir, file_path)
withopen(full_path, 'w') asf:
f.write('test content')

# Create symlinks
os.symlink(os.path.join(temp_dir, 'keep.py'),
os.path.join(temp_dir, 'ln-keep.py'))
os.symlink(os.path.join(temp_dir, 'dir/keep.py'),
os.path.join(temp_dir, 'ln-dir-keep.py'))
os.symlink(os.path.join(temp_dir, 'keep.py'),
os.path.join(temp_dir, 'dir/subdir/ln-keep.py'))

# Symlinks for folders
os.symlink(os.path.join(temp_dir, 'dir/subdir/ln-folder'),
os.path.join(temp_dir, 'ln-folder'))

# Create empty directories
os.makedirs(os.path.join(temp_dir, 'empty-folder'))

# Create skyignore file
skyignore_content="""
 # Current directory
 /remove.py
 /remove_dir
 /*.a
 /dir/*.b
 # Pattern match for all subdirectories
 *.sh
 remove.a
 """
skyignore_path=os.path.join(temp_dir, constants.SKY_IGNORE_FILE)
withopen(skyignore_path, 'w', encoding='utf-8') asf:
f.write(skyignore_content)

yieldtemp_dir