Skip to content

Commit 07fb12a

Browse files
ddurst-nvidiaddurst-nvidia
committed
chore(test): dedupe gateway gRPC router wiring in openshell-server integration tests
Signed-off-by: ddurst <267424412+ddurst-nvidia@users.noreply.github.com>
1 parent 31b8a05 commit 07fb12a

5 files changed

Lines changed: 52 additions & 96 deletions

File tree

‎crates/openshell-server/tests/auth_endpoint_integration.rs‎

Lines changed: 4 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ use serde::Deserialize;
2626
use std::net::SocketAddr;
2727
use tokio::net::TcpListener;
2828

29+
#[macro_use]
2930
mod common;
3031

3132
// ---------------------------------------------------------------------------
@@ -712,31 +713,16 @@ impl openshell_core::proto::open_shell_server::OpenShell for TestOpenShell {
712713
/// Uses `serve_connection_with_upgrades` to also support WebSocket upgrades.
713714
#[tokio::test]
714715
async fn plaintext_server_accepts_grpc_and_http() {
715-
use openshell_core::proto::{
716-
inference_server::InferenceServer, open_shell_server::OpenShellServer,
717-
};
718-
use openshell_server::{
719-
GatewayGrpcRouter, GatewayStandardHealth, MultiplexedService, OPENSHELL_SERVICE_NAME,
720-
health_router,
721-
};
716+
use openshell_server::{MultiplexedService, OPENSHELL_SERVICE_NAME, health_router};
722717
use tonic_health::pb::{
723718
HealthCheckRequest, health_check_response::ServingStatus, health_client::HealthClient,
724719
};
725720

726721
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
727722
let addr = listener.local_addr().unwrap();
728723

729-
let standard_health = GatewayStandardHealth::server(common::MAX_GRPC_DECODE);
730-
let reflection = tonic_reflection::server::Builder::configure()
731-
.register_encoded_file_descriptor_set(openshell_core::proto::FILE_DESCRIPTOR_SET)
732-
.register_encoded_file_descriptor_set(tonic_health::pb::FILE_DESCRIPTOR_SET)
733-
.build_v1()
734-
.unwrap();
735-
let openshell =
736-
OpenShellServer::new(TestOpenShell).max_decoding_message_size(common::MAX_GRPC_DECODE);
737-
let inference = InferenceServer::new(common::TestInference)
738-
.max_decoding_message_size(common::MAX_GRPC_DECODE);
739-
let grpc_service = GatewayGrpcRouter::new(standard_health, reflection, openshell, inference);
724+
let grpc_service =
725+
gateway_test_grpc_router!(common::openshell_max_decode_server(TestOpenShell));
740726
let http_service = health_router();
741727
let service = MultiplexedService::new(grpc_service, http_service);
742728

‎crates/openshell-server/tests/common/mod.rs‎

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,44 @@
66
use openshell_core::proto::{
77
GetClusterInferenceRequest, GetClusterInferenceResponse, GetInferenceBundleRequest,
88
GetInferenceBundleResponse, SetClusterInferenceRequest, SetClusterInferenceResponse,
9-
inference_server::Inference,
9+
inference_server::Inference, open_shell_server::OpenShellServer,
1010
};
1111
use tonic::{Response, Status};
1212

1313
/// Re-export the production gateway gRPC decode cap for test stacks.
1414
pub use openshell_server::MAX_GRPC_DECODE_SIZE as MAX_GRPC_DECODE;
1515

16+
/// Wrap an `OpenShell` impl with the same gRPC decode cap as production.
17+
#[must_use]
18+
pub fn openshell_max_decode_server<S>(inner: S) -> OpenShellServer<S> {
19+
OpenShellServer::new(inner).max_decoding_message_size(MAX_GRPC_DECODE)
20+
}
21+
22+
/// Build a [`GatewayGrpcRouter`](::openshell_server::GatewayGrpcRouter) matching production multiplex
23+
/// wiring: `grpc.health.v1`, reflection (`OpenShell` + `grpc.health` file descriptor sets), the given
24+
/// [`OpenShellServer`], and [`TestInference`].
25+
///
26+
/// Bring into scope with `#[macro_use] mod common;` in the integration test crate root.
27+
macro_rules! gateway_test_grpc_router {
28+
($openshell:expr) => {{
29+
::openshell_server::GatewayGrpcRouter::new(
30+
::openshell_server::GatewayStandardHealth::server(
31+
::openshell_server::MAX_GRPC_DECODE_SIZE,
32+
),
33+
::tonic_reflection::server::Builder::configure()
34+
.register_encoded_file_descriptor_set(::openshell_core::proto::FILE_DESCRIPTOR_SET)
35+
.register_encoded_file_descriptor_set(::tonic_health::pb::FILE_DESCRIPTOR_SET)
36+
.build_v1()
37+
.expect("OpenShell + grpc.health reflection descriptors"),
38+
$openshell,
39+
::openshell_core::proto::inference_server::InferenceServer::new(
40+
$crate::common::TestInference,
41+
)
42+
.max_decoding_message_size(::openshell_server::MAX_GRPC_DECODE_SIZE),
43+
)
44+
}};
45+
}
46+
1647
#[derive(Clone, Default)]
1748
pub struct TestInference;
1849

‎crates/openshell-server/tests/edge_tunnel_auth.rs‎

Lines changed: 5 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -44,18 +44,14 @@ use openshell_core::proto::{
4444
ListProvidersRequest, ListProvidersResponse, ListSandboxesRequest, ListSandboxesResponse,
4545
ProviderResponse, RevokeSshSessionRequest, RevokeSshSessionResponse, SandboxResponse,
4646
SandboxStreamEvent, ServiceStatus, SupervisorMessage, UpdateProviderRequest,
47-
WatchSandboxRequest,
48-
inference_server::InferenceServer,
49-
open_shell_server::{OpenShell, OpenShellServer},
50-
};
51-
use openshell_server::{
52-
GatewayGrpcRouter, GatewayStandardHealth, MultiplexedService, OPENSHELL_SERVICE_NAME,
53-
TlsAcceptor, health_router,
47+
WatchSandboxRequest, open_shell_server::OpenShell,
5448
};
49+
use openshell_server::{MultiplexedService, OPENSHELL_SERVICE_NAME, TlsAcceptor, health_router};
5550
use tonic_health::pb::{
5651
HealthCheckRequest, health_check_response::ServingStatus, health_client::HealthClient,
5752
};
5853

54+
#[macro_use]
5955
mod common;
6056
use rcgen::{CertificateParams, IsCa, KeyPair};
6157
use rustls::RootCertStore;
@@ -426,17 +422,8 @@ async fn start_test_server(
426422
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
427423
let addr = listener.local_addr().unwrap();
428424

429-
let standard_health = GatewayStandardHealth::server(common::MAX_GRPC_DECODE);
430-
let reflection = tonic_reflection::server::Builder::configure()
431-
.register_encoded_file_descriptor_set(openshell_core::proto::FILE_DESCRIPTOR_SET)
432-
.register_encoded_file_descriptor_set(tonic_health::pb::FILE_DESCRIPTOR_SET)
433-
.build_v1()
434-
.unwrap();
435-
let openshell =
436-
OpenShellServer::new(TestOpenShell).max_decoding_message_size(common::MAX_GRPC_DECODE);
437-
let inference = InferenceServer::new(common::TestInference)
438-
.max_decoding_message_size(common::MAX_GRPC_DECODE);
439-
let grpc_service = GatewayGrpcRouter::new(standard_health, reflection, openshell, inference);
425+
let grpc_service =
426+
gateway_test_grpc_router!(common::openshell_max_decode_server(TestOpenShell));
440427
let http_service = health_router();
441428
let service = MultiplexedService::new(grpc_service, http_service);
442429

‎crates/openshell-server/tests/multiplex_integration.rs‎

Lines changed: 6 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -18,16 +18,11 @@ use openshell_core::proto::{
1818
ListProvidersRequest, ListProvidersResponse, ListSandboxesRequest, ListSandboxesResponse,
1919
ProviderResponse, RevokeSshSessionRequest, RevokeSshSessionResponse, SandboxResponse,
2020
SandboxStreamEvent, ServiceStatus, SupervisorMessage, UpdateProviderRequest,
21-
WatchSandboxRequest,
22-
inference_server::InferenceServer,
23-
open_shell_client::OpenShellClient,
24-
open_shell_server::{OpenShell, OpenShellServer},
25-
};
26-
use openshell_server::{
27-
GatewayGrpcRouter, GatewayStandardHealth, MultiplexedService, OPENSHELL_SERVICE_NAME,
28-
health_router,
21+
WatchSandboxRequest, open_shell_client::OpenShellClient, open_shell_server::OpenShell,
2922
};
23+
use openshell_server::{MultiplexedService, OPENSHELL_SERVICE_NAME, health_router};
3024

25+
#[macro_use]
3126
mod common;
3227
use tokio::net::TcpListener;
3328
use tokio::sync::mpsc;
@@ -328,17 +323,7 @@ async fn serves_grpc_and_http_on_same_port() {
328323
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
329324
let addr = listener.local_addr().unwrap();
330325

331-
let standard_health = GatewayStandardHealth::server(common::MAX_GRPC_DECODE);
332-
let reflection = tonic_reflection::server::Builder::configure()
333-
.register_encoded_file_descriptor_set(openshell_core::proto::FILE_DESCRIPTOR_SET)
334-
.register_encoded_file_descriptor_set(tonic_health::pb::FILE_DESCRIPTOR_SET)
335-
.build_v1()
336-
.unwrap();
337-
let openshell =
338-
OpenShellServer::new(TestOpenShell).max_decoding_message_size(common::MAX_GRPC_DECODE);
339-
let inference = InferenceServer::new(common::TestInference)
340-
.max_decoding_message_size(common::MAX_GRPC_DECODE);
341-
let grpc_router = GatewayGrpcRouter::new(standard_health, reflection, openshell, inference);
326+
let grpc_router = gateway_test_grpc_router!(common::openshell_max_decode_server(TestOpenShell));
342327
let http_service = health_router();
343328
let service = MultiplexedService::new(grpc_router, http_service);
344329

@@ -418,17 +403,7 @@ async fn grpc_response_propagates_request_id() {
418403
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
419404
let addr = listener.local_addr().unwrap();
420405

421-
let standard_health = GatewayStandardHealth::server(common::MAX_GRPC_DECODE);
422-
let reflection = tonic_reflection::server::Builder::configure()
423-
.register_encoded_file_descriptor_set(openshell_core::proto::FILE_DESCRIPTOR_SET)
424-
.register_encoded_file_descriptor_set(tonic_health::pb::FILE_DESCRIPTOR_SET)
425-
.build_v1()
426-
.unwrap();
427-
let openshell =
428-
OpenShellServer::new(TestOpenShell).max_decoding_message_size(common::MAX_GRPC_DECODE);
429-
let inference = InferenceServer::new(common::TestInference)
430-
.max_decoding_message_size(common::MAX_GRPC_DECODE);
431-
let grpc_router = GatewayGrpcRouter::new(standard_health, reflection, openshell, inference);
406+
let grpc_router = gateway_test_grpc_router!(common::openshell_max_decode_server(TestOpenShell));
432407

433408
let x_request_id = http::HeaderName::from_static("x-request-id");
434409
let grpc_service = ServiceBuilder::new()
@@ -484,17 +459,7 @@ async fn standard_grpc_health_reflection_multiplexed() {
484459
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
485460
let addr = listener.local_addr().unwrap();
486461

487-
let standard_health = GatewayStandardHealth::server(common::MAX_GRPC_DECODE);
488-
let reflection = tonic_reflection::server::Builder::configure()
489-
.register_encoded_file_descriptor_set(openshell_core::proto::FILE_DESCRIPTOR_SET)
490-
.register_encoded_file_descriptor_set(tonic_health::pb::FILE_DESCRIPTOR_SET)
491-
.build_v1()
492-
.unwrap();
493-
let openshell =
494-
OpenShellServer::new(TestOpenShell).max_decoding_message_size(common::MAX_GRPC_DECODE);
495-
let inference = InferenceServer::new(common::TestInference)
496-
.max_decoding_message_size(common::MAX_GRPC_DECODE);
497-
let grpc_router = GatewayGrpcRouter::new(standard_health, reflection, openshell, inference);
462+
let grpc_router = gateway_test_grpc_router!(common::openshell_max_decode_server(TestOpenShell));
498463
let http_service = health_router();
499464
let service = MultiplexedService::new(grpc_router, http_service);
500465

‎crates/openshell-server/tests/multiplex_tls_integration.rs‎

Lines changed: 5 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -20,18 +20,14 @@ use openshell_core::proto::{
2020
ListProvidersRequest, ListProvidersResponse, ListSandboxesRequest, ListSandboxesResponse,
2121
ProviderResponse, RevokeSshSessionRequest, RevokeSshSessionResponse, SandboxResponse,
2222
SandboxStreamEvent, ServiceStatus, SupervisorMessage, UpdateProviderRequest,
23-
WatchSandboxRequest,
24-
inference_server::InferenceServer,
25-
open_shell_server::{OpenShell, OpenShellServer},
26-
};
27-
use openshell_server::{
28-
GatewayGrpcRouter, GatewayStandardHealth, MultiplexedService, OPENSHELL_SERVICE_NAME,
29-
TlsAcceptor, health_router,
23+
WatchSandboxRequest, open_shell_server::OpenShell,
3024
};
25+
use openshell_server::{MultiplexedService, OPENSHELL_SERVICE_NAME, TlsAcceptor, health_router};
3126
use tonic_health::pb::{
3227
HealthCheckRequest, health_check_response::ServingStatus, health_client::HealthClient,
3328
};
3429

30+
#[macro_use]
3531
mod common;
3632
use rcgen::{CertificateParams, IsCa, KeyPair};
3733
use rustls::RootCertStore;
@@ -406,17 +402,8 @@ async fn start_test_server(
406402
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
407403
let addr = listener.local_addr().unwrap();
408404

409-
let standard_health = GatewayStandardHealth::server(common::MAX_GRPC_DECODE);
410-
let reflection = tonic_reflection::server::Builder::configure()
411-
.register_encoded_file_descriptor_set(openshell_core::proto::FILE_DESCRIPTOR_SET)
412-
.register_encoded_file_descriptor_set(tonic_health::pb::FILE_DESCRIPTOR_SET)
413-
.build_v1()
414-
.unwrap();
415-
let openshell =
416-
OpenShellServer::new(TestOpenShell).max_decoding_message_size(common::MAX_GRPC_DECODE);
417-
let inference = InferenceServer::new(common::TestInference)
418-
.max_decoding_message_size(common::MAX_GRPC_DECODE);
419-
let grpc_service = GatewayGrpcRouter::new(standard_health, reflection, openshell, inference);
405+
let grpc_service =
406+
gateway_test_grpc_router!(common::openshell_max_decode_server(TestOpenShell));
420407
let http_service = health_router();
421408
let service = MultiplexedService::new(grpc_service, http_service);
422409

0 commit comments

Comments
 (0)