Line data    Source code

       1              : use crate::proto::dashql_test::{
       2              :     test_service_server::{TestService, TestServiceServer}, TestServerStreamingRequest, TestServerStreamingResponse, TestUnaryRequest, TestUnaryResponse
       3              : };
       4              : use std::{net::SocketAddr, pin::Pin, sync::{Arc, Mutex}};
       5              : use tokio::sync::oneshot;
       6              : use tokio_stream::{wrappers::ReceiverStream, Stream};
       7              : use tonic::{metadata::MetadataMap, Status};
       8              : 
       9              : pub type UnaryResponseSender = tokio::sync::mpsc::Sender<Result<TestUnaryResponse, tonic::Status>>;
      10              : pub type ServerStreamingResponseSender = tokio::sync::mpsc::Sender<Result<TestServerStreamingResponse, tonic::Status>>;
      11              : 
      12              : /// A dashql.test.TestService that hands every call to the test for a response
      13              : pub struct GrpcServiceMock {
      14              :     /// Receives the parameters of a unary call together with the sender of its response
      15              :     pub setup_unary: tokio::sync::mpsc::Sender<(TestUnaryRequest, UnaryResponseSender)>,
      16              :     /// Receives the parameters of a streaming call together with the sender of its responses
      17              :     pub setup_server_streaming: tokio::sync::mpsc::Sender<(TestServerStreamingRequest, ServerStreamingResponseSender)>,
      18              :     /// The metadata of the received calls
      19              :     pub received_metadata: Arc<Mutex<Vec<MetadataMap>>>,
      20              : }
      21              : 
      22              : impl GrpcServiceMock {
      23           13 :     pub fn new() -> (
      24           13 :         Self,
      25           13 :         tokio::sync::mpsc::Receiver<(TestUnaryRequest, UnaryResponseSender)>,
      26           13 :         tokio::sync::mpsc::Receiver<(TestServerStreamingRequest, ServerStreamingResponseSender)>
      27           13 :     ) {
      28           13 :         let (send_unary_setup, recv_unary_setup) = tokio::sync::mpsc::channel(1);
      29           13 :         let (send_server_streaming_setup, recv_server_streaming_setup) = tokio::sync::mpsc::channel(1);
      30           13 :         (Self { setup_unary: send_unary_setup, setup_server_streaming: send_server_streaming_setup, received_metadata: Default::default() }, recv_unary_setup, recv_server_streaming_setup)
      31           13 :     }
      32              : }
      33              : 
      34              : type ServerStreamingResponseStream = Pin<Box<dyn Stream<Item = Result<TestServerStreamingResponse, Status>> + Send>>;
      35              : 
      36              : #[tonic::async_trait]
      37              : impl TestService for GrpcServiceMock {
      38              :     type TestServerStreamingStream = ServerStreamingResponseStream;
      39              : 
      40              :     async fn test_unary(
      41              :         &self,
      42              :         request: tonic::Request<TestUnaryRequest>
      43            4 :     ) -> Result<tonic::Response<TestUnaryResponse>, tonic::Status> {
      44              :         self.received_metadata.lock().unwrap().push(request.metadata().clone());
      45              :         // Setup a channel for sending the result
      46              :         let (result_sender, mut receiver) = tokio::sync::mpsc::channel(1);
      47              :         let test_setup = self.setup_unary.clone();
      48              : 
      49              :         // Pass the result_sender back to the test, together with the request params
      50              :         let params = request.into_inner();
      51              :         test_setup.send((params, result_sender)).await.unwrap();
      52              : 
      53              :         // Await the response
      54              :         let response = receiver.recv().await.unwrap()?;
      55              :         Ok(tonic::Response::new(response))
      56            4 :     }
      57              : 
      58              :     async fn test_server_streaming(
      59              :         &self,
      60              :         request: tonic::Request<TestServerStreamingRequest>,
      61            6 :     ) -> Result<tonic::Response<Self::TestServerStreamingStream>, tonic::Status> {
      62              :         self.received_metadata.lock().unwrap().push(request.metadata().clone());
      63              :         // Setup a channel for sending results
      64              :         let (result_sender, receive) = tokio::sync::mpsc::channel(10);
      65              :         let test_setup = self.setup_server_streaming.clone();
      66              : 
      67              :         // Pass the result_sender back to the test, together with the request params
      68              :         let params = request.into_inner();
      69              :         test_setup.send((params, result_sender)).await.unwrap();
      70              : 
      71              :         // Wire up a receiver stream to the gRPC output
      72              :         let out = ReceiverStream::new(receive);
      73              :         Ok(tonic::Response::new(
      74              :             Box::pin(out) as Self::TestServerStreamingStream
      75              :         ))
      76            6 :     }
      77              : }
      78              : 
      79           13 : pub async fn spawn_grpc_test_service_mock(mock: GrpcServiceMock) -> (SocketAddr, oneshot::Sender<()>) {
      80           13 :     let service = TestServiceServer::new(mock);
      81              : 
      82              :     // Bind to port `0` so the OS finds a free port
      83           13 :     let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
      84           13 :     let addr = listener.local_addr().unwrap();
      85              : 
      86           13 :     let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
      87              : 
      88           13 :     tokio::spawn(async move {
      89           13 :         tonic::transport::Server::builder()
      90           13 :             .add_service(service)
      91           13 :             .serve_with_incoming_shutdown(
      92           13 :                 tokio_stream::wrappers::TcpListenerStream::new(listener),
      93           13 :                 async { drop(shutdown_rx.await) },
      94              :             )
      95           13 :             .await
      96            0 :             .unwrap();
      97            0 :     });
      98           13 :     (addr, shutdown_tx)
      99           13 : }