wehub-resource-sync
61 lines
1.8 KiB
Python
61 lines
1.8 KiB
Python
import sys
|
|
|
|
import grpc
|
|
import pytest
|
|
|
|
from ray.util.client.worker import Worker
|
|
|
|
|
|
class Credentials(grpc.ChannelCredentials):
|
|
def __init__(self, name):
|
|
self.name = name
|
|
|
|
|
|
def test_grpc_client_credentials_are_passed_to_channel(monkeypatch):
|
|
class Stop(Exception):
|
|
def __init__(self, credentials):
|
|
self.credentials = credentials
|
|
|
|
class MockChannel:
|
|
def __init__(self, conn_str, credentials, options, compression):
|
|
self.credentials = credentials
|
|
|
|
def subscribe(self, f):
|
|
raise Stop(self.credentials)
|
|
|
|
def mock_secure_channel(conn_str, credentials, options=None, compression=None):
|
|
return MockChannel(conn_str, credentials, options, compression)
|
|
|
|
monkeypatch.setattr(grpc, "secure_channel", mock_secure_channel)
|
|
|
|
# Credentials should be respected whether secure is set or not.
|
|
|
|
with pytest.raises(Stop) as stop:
|
|
Worker(secure=False, _credentials=Credentials("test"))
|
|
assert stop.value.credentials.name == "test"
|
|
|
|
with pytest.raises(Stop) as stop:
|
|
Worker(secure=True, _credentials=Credentials("test"))
|
|
assert stop.value.credentials.name == "test"
|
|
|
|
|
|
def test_grpc_client_credentials_are_generated(monkeypatch):
|
|
# Test that credentials are generated when secure is True, but _credentials
|
|
# isn't passed.
|
|
class Stop(Exception):
|
|
def __init__(self, result):
|
|
self.result = result
|
|
|
|
def mock_gen_credentials():
|
|
raise Stop("ssl_channel_credentials called")
|
|
|
|
monkeypatch.setattr(grpc, "ssl_channel_credentials", mock_gen_credentials)
|
|
|
|
with pytest.raises(Stop) as stop:
|
|
Worker(secure=True)
|
|
assert stop.value.result == "ssl_channel_credentials called"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|