def test_load_via_env_var():
port = find_free_port()
python_file = file_relative_path(__file__, "grpc_repo.py")
subprocess_args = [
"dagster",
"api",
"grpc",
"--python-file",
python_file,
]
with environ({
"DAGSTER_CLI_API_GRPC_HOST": "localhost",
"DAGSTER_CLI_API_GRPC_PORT": str(port)
}):
process = subprocess.Popen(
subprocess_args,
stdout=subprocess.PIPE,
)
try:
wait_for_grpc_server(
process, DagsterGrpcClient(port=port, host="localhost"),
subprocess_args)
assert DagsterGrpcClient(port=port).ping("foobar") == "foobar"
finally:
process.terminate()
def test_load_with_empty_working_directory(capfd):
port = find_free_port()
# File that will fail if working directory isn't set to default
python_file = file_relative_path(__file__, "grpc_repo_with_local_import.py")
subprocess_args = [
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
]
with new_cwd(os.path.dirname(__file__)):
process = subprocess.Popen(
subprocess_args,
stdout=subprocess.PIPE,
)
try:
wait_for_grpc_server(
process, DagsterGrpcClient(port=port, host="localhost"), subprocess_args
)
assert DagsterGrpcClient(port=port).ping("foobar") == "foobar"
finally:
process.terminate()
# indicating the working directory is empty fails
port = find_free_port()
subprocess_args = [
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
"--empty-working-directory",
]
process = subprocess.Popen(
subprocess_args,
stdout=subprocess.PIPE,
)
try:
with pytest.raises(Exception):
wait_for_grpc_server(
process, DagsterGrpcClient(port=port, host="localhost"), subprocess_args
)
process.wait()
_, err = capfd.readouterr()
assert "No module named" in err
finally:
if process.poll() is None:
process.terminate()
def test_lazy_load_with_error():
port = find_free_port()
python_file = file_relative_path(__file__, "grpc_repo_with_error.py")
subprocess_args = [
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
"--lazy-load-user-code",
]
process = subprocess.Popen(subprocess_args, stdout=subprocess.PIPE)
try:
wait_for_grpc_server(process,
DagsterGrpcClient(port=port, host="localhost"),
subprocess_args)
list_repositories_response = deserialize_json_to_dagster_namedtuple(
DagsterGrpcClient(port=port).list_repositories())
assert isinstance(list_repositories_response, SerializableErrorInfo)
assert "No module named" in list_repositories_response.message
finally:
process.terminate()
def test_lazy_load_via_env_var():
with environ({"DAGSTER_CLI_API_GRPC_LAZY_LOAD_USER_CODE": "1"}):
port = find_free_port()
python_file = file_relative_path(__file__, "grpc_repo_with_error.py")
subprocess_args = [
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
]
process = subprocess.Popen(
subprocess_args,
stdout=subprocess.PIPE,
)
try:
wait_for_grpc_server(
process, DagsterGrpcClient(port=port, host="localhost"),
subprocess_args)
list_repositories_response = deserialize_json_to_dagster_namedtuple(
DagsterGrpcClient(port=port).list_repositories())
assert isinstance(list_repositories_response,
SerializableErrorInfo)
assert "No module named" in list_repositories_response.message
finally:
process.terminate()
def test_streaming():
port = find_free_port()
python_file = file_relative_path(__file__, "grpc_repo.py")
ipc_output_file = _get_ipc_output_file()
process = subprocess.Popen(
[
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
"--ipc-output-file",
ipc_output_file,
],
stdout=subprocess.PIPE,
)
try:
wait_for_grpc_server(ipc_output_file)
api_client = DagsterGrpcClient(port=port)
results = [
result for result in api_client.streaming_ping(sequence_length=10,
echo="foo")
]
assert len(results) == 10
for sequence_number, result in enumerate(results):
assert result["sequence_number"] == sequence_number
assert result["echo"] == "foo"
finally:
process.terminate()
def __init__(self, origin):
from dagster.grpc.client import DagsterGrpcClient
from dagster.grpc.server_watcher import create_grpc_watch_thread
self._origin = check.inst_param(origin, "origin",
GrpcServerRepositoryLocationOrigin)
port = self.origin.port
socket = self.origin.socket
host = self.origin.host
self._watch_thread_shutdown_event = None
self._watch_thread = None
try:
self.client = DagsterGrpcClient(port=port,
socket=socket,
host=host)
list_repositories_response = sync_list_repositories_grpc(
self.client)
self.server_id = sync_get_server_id(self.client)
self.repository_names = set(
symbol.repository_name
for symbol in list_repositories_response.repository_symbols)
self._state_subscribers = []
self._watch_thread_shutdown_event, self._watch_thread = create_grpc_watch_thread(
self.client,
on_updated=lambda new_server_id: self.
_send_state_event_to_subscribers(
LocationStateChangeEvent(
LocationStateChangeEventType.LOCATION_UPDATED,
location_name=self.location_name,
message="Server has been updated.",
server_id=new_server_id,
)),
on_error=lambda: self._send_state_event_to_subscribers(
LocationStateChangeEvent(
LocationStateChangeEventType.LOCATION_ERROR,
location_name=self.location_name,
message=
"Unable to reconnect to server. You can reload the server once it is "
"reachable again",
)),
)
self._watch_thread.start()
self.executable_path = list_repositories_response.executable_path
self.repository_code_pointer_dict = (
list_repositories_response.repository_code_pointer_dict)
self.container_image = self._reload_current_image()
except:
self.cleanup()
raise
def test_cleanup_after_force_terminate(run_config):
with instance_for_test() as instance, get_managed_grpc_server_workspace(
instance) as workspace:
external_pipeline = (
workspace.get_repository_location("test").get_repository(
"nope").get_full_external_pipeline("sleepy_pipeline"))
pipeline_run = instance.create_run_for_pipeline(
pipeline_def=sleepy_pipeline,
run_config=run_config,
external_pipeline_origin=external_pipeline.get_external_origin(),
pipeline_code_origin=external_pipeline.get_python_origin(),
)
run_id = pipeline_run.run_id
instance.launch_run(pipeline_run.run_id, workspace)
poll_for_step_start(instance, run_id)
# simulate the sequence of events that happen during force-termination:
# run moves immediately into canceled status while termination happens
instance.report_run_canceling(pipeline_run)
instance.report_run_canceled(pipeline_run)
reloaded_run = instance.get_run_by_id(run_id)
grpc_info = json.loads(reloaded_run.tags.get(GRPC_INFO_TAG))
client = DagsterGrpcClient(
port=grpc_info.get("port"),
socket=grpc_info.get("socket"),
host=grpc_info.get("host"),
)
client.cancel_execution(CancelExecutionRequest(run_id=run_id))
# Wait for the run worker to clean up
start_time = time.time()
while True:
if time.time() - start_time > 30:
raise Exception("Timed out waiting for cleanup message")
logs = instance.all_logs(run_id)
if any([
"Computational resources were cleaned up after the run was forcibly marked as canceled."
in str(event) for event in logs
]):
break
time.sleep(1)
assert instance.get_run_by_id(
run_id).status == PipelineRunStatus.CANCELED
def open_server_process(
port,
socket,
loadable_target_origin=None,
max_workers=None,
heartbeat=False,
heartbeat_timeout=30,
fixed_server_id=None,
startup_timeout=20,
):
check.invariant((port or socket) and not (port and socket), "Set only port or socket")
check.opt_inst_param(loadable_target_origin, "loadable_target_origin", LoadableTargetOrigin)
check.opt_int_param(max_workers, "max_workers")
from dagster.core.test_utils import get_mocked_system_timezone
mocked_system_timezone = get_mocked_system_timezone()
executable_path = loadable_target_origin.executable_path if loadable_target_origin else None
subprocess_args = (
(
get_python_environment_entry_point(executable_path)
if executable_path
else DEFAULT_DAGSTER_ENTRY_POINT
)
+ ["api", "grpc"]
+ ["--lazy-load-user-code"]
+ (["--port", str(port)] if port else [])
+ (["--socket", socket] if socket else [])
+ (["-n", str(max_workers)] if max_workers else [])
+ (["--heartbeat"] if heartbeat else [])
+ (["--heartbeat-timeout", str(heartbeat_timeout)] if heartbeat_timeout else [])
+ (["--fixed-server-id", fixed_server_id] if fixed_server_id else [])
+ (["--override-system-timezone", mocked_system_timezone] if mocked_system_timezone else [])
+ (["--log-level", "WARNING"]) # don't log INFO messages for automatically spun up servers
+ (["--use-python-environment-entry-point"] if executable_path else [])
)
if loadable_target_origin:
subprocess_args += loadable_target_origin.get_cli_args()
server_process = open_ipc_subprocess(subprocess_args)
from dagster.grpc.client import DagsterGrpcClient
client = DagsterGrpcClient(
port=port,
socket=socket,
host="localhost",
)
try:
wait_for_grpc_server(server_process, client, subprocess_args, timeout=startup_timeout)
except:
if server_process.poll() is None:
server_process.terminate()
raise
return server_process
def test_lazy_load_via_env_var():
with environ({"DAGSTER_CLI_API_GRPC_LAZY_LOAD_USER_CODE": "1"}):
port = find_free_port()
python_file = file_relative_path(__file__, "grpc_repo_with_error.py")
ipc_output_file = _get_ipc_output_file()
process = subprocess.Popen(
[
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
"--ipc-output-file",
ipc_output_file,
],
stdout=subprocess.PIPE,
)
try:
wait_for_grpc_server(process, ipc_output_file)
list_repositories_response = DagsterGrpcClient(port=port).list_repositories()
assert isinstance(list_repositories_response, SerializableErrorInfo)
assert "No module named" in list_repositories_response.message
finally:
process.terminate()
def test_load_with_empty_working_directory(capfd):
port = find_free_port()
# File that will fail if working directory isn't set to default
python_file = file_relative_path(__file__, "grpc_repo_with_local_import.py")
with new_cwd(os.path.dirname(__file__)):
ipc_output_file = _get_ipc_output_file()
process = subprocess.Popen(
[
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
"--ipc-output-file",
ipc_output_file,
],
stdout=subprocess.PIPE,
)
try:
wait_for_grpc_server(process, ipc_output_file)
assert DagsterGrpcClient(port=port).ping("foobar") == "foobar"
finally:
process.terminate()
# indicating the working directory is empty fails
ipc_output_file = _get_ipc_output_file()
process = subprocess.Popen(
[
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
"--empty-working-directory",
"--ipc-output-file",
ipc_output_file,
],
stdout=subprocess.PIPE,
)
try:
with pytest.raises(DagsterUserCodeProcessError):
wait_for_grpc_server(process, ipc_output_file)
process.wait()
_, err = capfd.readouterr()
assert "No module named" in err
finally:
if process.poll() is None:
process.terminate()
def test_load_with_invalid_param(capfd):
port = find_free_port()
python_file = file_relative_path(__file__, "grpc_repo.py")
subprocess_args = [
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
"--foo-param",
"bar_value",
]
process = subprocess.Popen(
subprocess_args,
stdout=subprocess.PIPE,
)
try:
with pytest.raises(
Exception,
match='gRPC server exited with return code 2 while starting up with the command: "dagster api grpc --port',
):
wait_for_grpc_server(
process, DagsterGrpcClient(port=port, host="localhost"), subprocess_args
)
finally:
process.terminate()
_, err = capfd.readouterr()
assert "no such option" in err
def test_load_with_error(capfd):
port = find_free_port()
python_file = file_relative_path(__file__, "grpc_repo_with_error.py")
subprocess_args = [
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
]
process = subprocess.Popen(
subprocess_args,
stdout=subprocess.PIPE,
)
try:
with pytest.raises(Exception):
wait_for_grpc_server(
process, DagsterGrpcClient(port=port, host="localhost"),
subprocess_args)
process.wait()
_, err = capfd.readouterr()
assert "No module named" in err
finally:
if process.poll() is None:
process.terminate()
def test_load_grpc_server_python_env():
port = find_free_port()
python_file = file_relative_path(__file__, "grpc_repo.py")
subprocess_args = [
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
"--use-python-environment-entry-point",
]
process = subprocess.Popen(subprocess_args)
try:
client = DagsterGrpcClient(port=port, host="localhost")
wait_for_grpc_server(process, client, subprocess_args)
list_repositories_response = sync_list_repositories_grpc(client)
assert list_repositories_response.entry_point == [
sys.executable, "-m", "dagster"
]
assert list_repositories_response.executable_path == sys.executable
finally:
process.terminate()
def test_crash_during_load():
port = find_free_port()
python_file = file_relative_path(__file__, "crashy_grpc_repo.py")
subprocess_args = [
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
]
process = subprocess.Popen(
subprocess_args,
stdout=subprocess.PIPE,
)
try:
with pytest.raises(
Exception,
match=re.escape(
'gRPC server exited with return code 123 while starting up with the command: "dagster api grpc --port'
),
):
wait_for_grpc_server(
process, DagsterGrpcClient(port=port, host="localhost"),
subprocess_args)
finally:
if process.poll() is None:
process.terminate()
def docker_grpc_client(dagster_docker_image, grpc_host, grpc_port): # pylint: disable=redefined-outer-name, unused-argument
if not IS_BUILDKITE:
docker_service_up(file_relative_path(__file__, "docker-compose.yml"),
"dagster-grpc-server")
wait_for_connection(grpc_host, grpc_port)
yield DagsterGrpcClient(port=grpc_port, host=grpc_host)
def test_ping():
port = find_free_port()
python_file = file_relative_path(__file__, "grpc_repo.py")
ipc_output_file = _get_ipc_output_file()
process = subprocess.Popen(
[
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
"--ipc-output-file",
ipc_output_file,
],
stdout=subprocess.PIPE,
)
try:
wait_for_grpc_server(ipc_output_file)
assert DagsterGrpcClient(port=port).ping("foobar") == "foobar"
finally:
process.terminate()
def test_grpc_watch_thread_server_update():
port = find_free_port()
called = {}
def on_updated():
called["yup"] = True
# Create initial server
server_process = open_server_process(port=port, socket=None)
try:
# Start watch thread
client = DagsterGrpcClient(port=port)
watch_interval = 4
shutdown_event, watch_thread = create_grpc_watch_thread(
client, on_updated=on_updated, watch_interval=watch_interval)
watch_thread.start()
time.sleep(watch_interval * 2)
finally:
interrupt_ipc_subprocess_pid(server_process.pid)
assert not called
# Create updated server
server_process = open_server_process(port=port, socket=None)
try:
time.sleep(watch_interval * 2)
finally:
interrupt_ipc_subprocess_pid(server_process.pid)
shutdown_event.set()
watch_thread.join()
assert called
def test_run_grpc_watch_thread():
client = DagsterGrpcClient(port=8080)
shutdown_event, watch_thread = create_grpc_watch_thread(client)
watch_thread.start()
shutdown_event.set()
watch_thread.join()
def test_lazy_load_with_error():
port = find_free_port()
python_file = file_relative_path(__file__, "grpc_repo_with_error.py")
ipc_output_file = _get_ipc_output_file()
process = subprocess.Popen(
[
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
"--lazy-load-user-code",
"--ipc-output-file",
ipc_output_file,
],
stdout=subprocess.PIPE,
)
try:
wait_for_grpc_server(ipc_output_file)
list_repositories_response = DagsterGrpcClient(
port=port).list_repositories()
assert isinstance(list_repositories_response, SerializableErrorInfo)
assert "No module named" in list_repositories_response.message
finally:
process.terminate()
def create_grpc_server_location(port, socket, host, location_name=None):
from dagster.grpc.client import DagsterGrpcClient
check.opt_int_param(port, 'port')
check.opt_str_param(socket, 'socket')
check.str_param(host, 'host')
check.opt_str_param(location_name, 'location_name')
client = DagsterGrpcClient(port=port, socket=socket, host=host)
list_repositories_response = sync_list_repositories_grpc(client)
repository_names = set(
symbol.repository_name
for symbol in list_repositories_response.repository_symbols)
return GrpcServerRepositoryLocationHandle(
port=port,
socket=socket,
host=host,
location_name=location_name if location_name else
_assign_grpc_location_name(port, socket, host),
client=client,
repository_names=repository_names,
)
def test_sensor_timeout():
port = find_free_port()
python_file = file_relative_path(__file__, "grpc_repo.py")
subprocess_args = [
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
]
process = subprocess.Popen(
subprocess_args,
stdout=subprocess.PIPE,
)
try:
wait_for_grpc_server(
process, DagsterGrpcClient(port=port, host="localhost"), subprocess_args
)
client = DagsterGrpcClient(port=port)
with instance_for_test() as instance:
repo_origin = ExternalRepositoryOrigin(
repository_location_origin=GrpcServerRepositoryLocationOrigin(
port=port, host="localhost"
),
repository_name="bar_repo",
)
with pytest.raises(DagsterUserCodeUnreachableError) as exc_info:
client.external_sensor_execution(
sensor_execution_args=SensorExecutionArgs(
repository_origin=repo_origin,
instance_ref=instance.get_ref(),
sensor_name="slow_sensor",
last_completion_time=None,
last_run_key=None,
cursor=None,
),
timeout=2,
)
assert "Deadline Exceeded" in str(exc_info.getrepr())
# Call succeeds without the timeout
client.external_sensor_execution(
sensor_execution_args=SensorExecutionArgs(
repository_origin=repo_origin,
instance_ref=instance.get_ref(),
sensor_name="slow_sensor",
last_completion_time=None,
last_run_key=None,
cursor=None,
),
)
finally:
process.terminate()
def open_server_process(
port,
socket,
loadable_target_origin=None,
max_workers=None,
heartbeat=False,
heartbeat_timeout=30,
fixed_server_id=None,
):
check.invariant((port or socket) and not (port and socket),
"Set only port or socket")
check.opt_inst_param(loadable_target_origin, "loadable_target_origin",
LoadableTargetOrigin)
check.opt_int_param(max_workers, "max_workers")
from dagster.core.test_utils import get_mocked_system_timezone
mocked_system_timezone = get_mocked_system_timezone()
subprocess_args = (
[
loadable_target_origin.executable_path if loadable_target_origin
and loadable_target_origin.executable_path else sys.executable,
"-m",
"dagster.grpc",
] + ["--lazy-load-user-code"] +
(["--port", str(port)] if port else []) +
(["--socket", socket] if socket else []) +
(["-n", str(max_workers)] if max_workers else []) +
(["--heartbeat"] if heartbeat else []) +
(["--heartbeat-timeout", str(heartbeat_timeout)]
if heartbeat_timeout else []) +
(["--fixed-server-id", fixed_server_id] if fixed_server_id else []) +
(["--override-system-timezone", mocked_system_timezone]
if mocked_system_timezone else []))
if loadable_target_origin:
subprocess_args += loadable_target_origin.get_cli_args()
server_process = open_ipc_subprocess(subprocess_args)
from dagster.grpc.client import DagsterGrpcClient
client = DagsterGrpcClient(
port=port,
socket=socket,
host="localhost",
)
try:
wait_for_grpc_server(server_process, client, subprocess_args)
except:
if server_process.poll() is None:
server_process.terminate()
raise
return server_process
def create_client(self):
from dagster.grpc.client import DagsterGrpcClient
return DagsterGrpcClient(
port=self.port,
socket=self.socket,
host=self.host,
use_ssl=bool(self.use_ssl),
)
def test_load_grpc_server(capfd):
port = find_free_port()
python_file = file_relative_path(__file__, "grpc_repo.py")
subprocess_args = [
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
]
process = subprocess.Popen(subprocess_args)
try:
client = DagsterGrpcClient(port=port, host="localhost")
wait_for_grpc_server(process, client, subprocess_args)
assert client.ping("foobar") == "foobar"
list_repositories_response = sync_list_repositories_grpc(client)
assert list_repositories_response.entry_point == ["dagster"]
assert list_repositories_response.executable_path == sys.executable
subprocess.check_call(
["dagster", "api", "grpc-health-check", "--port",
str(port)])
ssl_result = subprocess.run( # pylint:disable=subprocess-run-check
[
"dagster", "api", "grpc-health-check", "--port",
str(port), "--use-ssl"
])
assert ssl_result.returncode == 1
finally:
process.terminate()
out, _err = capfd.readouterr()
assert f"Started Dagster code server for file {python_file} on port {port} in process" in out
def test_load_with_container_context(capfd):
port = find_free_port()
python_file = file_relative_path(__file__, "grpc_repo.py")
container_context = {
"k8s": {
"image_pull_policy": "Never",
"image_pull_secrets": [{"name": "your_secret"}],
}
}
subprocess_args = [
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
"--container-context",
json.dumps(container_context),
]
process = subprocess.Popen(subprocess_args)
try:
client = DagsterGrpcClient(port=port, host="localhost")
wait_for_grpc_server(process, client, subprocess_args)
assert client.ping("foobar") == "foobar"
list_repositories_response = sync_list_repositories_grpc(client)
assert list_repositories_response.entry_point == ["dagster"]
assert list_repositories_response.executable_path == sys.executable
assert list_repositories_response.container_context == container_context
finally:
process.terminate()
out, _err = capfd.readouterr()
assert f"Started Dagster code server for file {python_file} on port {port} in process" in out
def __init__(self, origin):
from dagster.grpc.client import DagsterGrpcClient
self.origin = check.inst_param(origin, "origin",
GrpcServerRepositoryLocationOrigin)
port = self.origin.port
socket = self.origin.socket
host = self.origin.host
self.client = DagsterGrpcClient(port=port, socket=socket, host=host)
list_repositories_response = sync_list_repositories_grpc(self.client)
self.repository_names = set(
symbol.repository_name
for symbol in list_repositories_response.repository_symbols)
self.executable_path = list_repositories_response.executable_path
self.repository_code_pointer_dict = list_repositories_response.repository_code_pointer_dict
def test_grpc_watch_thread_server_complex_cycle_2():
# Server goes down, comes back up as the same server three times, then goes away and comes
# back as a new server
port = find_free_port()
fixed_server_id = "fixed_id"
events = []
called = {}
def on_disconnect():
events.append("on_disconnect")
def on_reconnected():
events.append("on_reconnected")
def on_updated(_):
events.append("on_updated")
def on_error():
called["on_error"] = True
events.append("on_error")
# Create initial server
open_server_process(port=port, socket=None, fixed_server_id=fixed_server_id)
# Start watch thread
client = DagsterGrpcClient(port=port)
watch_interval = 1 # This is a faster watch interval than we would use in practice
shutdown_event, watch_thread = create_grpc_watch_thread(
client,
on_disconnect=on_disconnect,
on_reconnected=on_reconnected,
on_updated=on_updated,
on_error=on_error,
watch_interval=watch_interval,
max_reconnect_attempts=3,
)
watch_thread.start()
time.sleep(watch_interval * 3)
cycles = 3
for x in range(1, cycles + 1):
# Simulate server restart three times with same server ID
client.shutdown_server()
wait_for_condition(lambda: events.count("on_disconnect") == x, watch_interval)
open_server_process(port=port, socket=None, fixed_server_id=fixed_server_id)
wait_for_condition(lambda: events.count("on_reconnected") == x, watch_interval)
# Simulate server failure
client.shutdown_server()
# Wait for reconnect attempts to exhaust and on_error callback to be called
wait_for_condition(lambda: called.get("on_error"), watch_interval)
shutdown_event.set()
watch_thread.join()
assert events[-1] == "on_error"
def test_grpc_watch_thread_server_error():
port = find_free_port()
fixed_server_id = "fixed_id"
called = {}
def on_disconnect():
called["on_disconnect"] = True
def on_error():
called["on_error"] = True
def should_not_be_called():
raise Exception("This method should not be called")
# Create initial server
server_process = open_server_process(port=port,
socket=None,
fixed_server_id=fixed_server_id)
# Start watch thread
client = DagsterGrpcClient(port=port)
watch_interval = 1
max_reconnect_attempts = 3
shutdown_event, watch_thread = create_grpc_watch_thread(
client,
on_disconnect=on_disconnect,
on_reconnected=should_not_be_called,
on_updated=should_not_be_called,
on_error=on_error,
watch_interval=watch_interval,
max_reconnect_attempts=max_reconnect_attempts,
)
watch_thread.start()
# Wait three seconds, simulate restart failure
time.sleep(watch_interval * 3)
interrupt_ipc_subprocess_pid(server_process.pid)
# Wait for reconnect attempts to exhaust and on_error callback to be called
start_time = time.time()
while not called.get("on_error"):
if time.time() - start_time > 30:
break
time.sleep(1)
shutdown_event.set()
watch_thread.join()
assert called["on_disconnect"]
assert called["on_error"]
def test_ping():
port = find_free_port()
python_file = file_relative_path(__file__, "grpc_repo.py")
subprocess_args = [
"dagster",
"api",
"grpc",
"--port",
str(port),
"--python-file",
python_file,
]
process = subprocess.Popen(subprocess_args, stdout=subprocess.PIPE)
try:
wait_for_grpc_server(
process, DagsterGrpcClient(port=port, host="localhost"), subprocess_args
)
assert DagsterGrpcClient(port=port).ping("foobar") == "foobar"
finally:
process.terminate()
def test_streaming():
port = find_free_port()
python_file = file_relative_path(__file__, 'grpc_repo.py')
process = subprocess.Popen(
[
'dagster', 'api', 'grpc', '--port',
str(port), '--python-file', python_file
],
stdout=subprocess.PIPE,
)
try:
wait_for_grpc_server(process)
api_client = DagsterGrpcClient(port=port)
results = [
result for result in api_client.streaming_ping(sequence_length=10,
echo='foo')
]
assert len(results) == 10
for sequence_number, result in enumerate(results):
assert result['sequence_number'] == sequence_number
assert result['echo'] == 'foo'
finally:
process.terminate()
