Line data    Source code

       1              : use crate::proto::salesforce_hyperdb_grpc_v1::{
       2              :     hyper_service_server::{HyperService, HyperServiceServer},
       3              :     QueryParam, QueryResult,
       4              : };
       5              : use std::{net::SocketAddr, pin::Pin};
       6              : use tokio::sync::oneshot;
       7              : use tokio_stream::{wrappers::ReceiverStream, Stream};
       8              : use tonic::Status;
       9              : 
      10              : pub type QueryResultSender = tokio::sync::mpsc::Sender<Result<QueryResult, tonic::Status>>;
      11              : 
      12              : /// A Hyper gRPC service that hands every query to the test for its results
      13              : pub struct HyperServiceMock {
      14              :     /// Receives the parameters of a query together with the sender of its results
      15              :     pub setup_execute_query: tokio::sync::mpsc::Sender<(QueryParam, QueryResultSender)>,
      16              : }
      17              : 
      18              : impl HyperServiceMock {
      19            2 :     pub fn new() -> (
      20            2 :         Self,
      21            2 :         tokio::sync::mpsc::Receiver<(QueryParam, QueryResultSender)>,
      22            2 :     ) {
      23            2 :         let (send, recv) = tokio::sync::mpsc::channel(1);
      24            2 :         (Self { setup_execute_query: send }, recv)
      25            2 :     }
      26              : }
      27              : 
      28              : type ExecuteQueryResponseStream = Pin<Box<dyn Stream<Item = Result<QueryResult, Status>> + Send>>;
      29              : 
      30              : #[tonic::async_trait]
      31              : impl HyperService for HyperServiceMock {
      32              :     type ExecuteQueryStream = ExecuteQueryResponseStream;
      33              : 
      34              :     async fn execute_query(
      35              :         &self,
      36              :         request: tonic::Request<QueryParam>,
      37            2 :     ) -> Result<tonic::Response<Self::ExecuteQueryStream>, tonic::Status> {
      38              :         // Setup a channel for sending results
      39              :         let (result_sender, receive) = tokio::sync::mpsc::channel(10);
      40              :         let test_setup = self.setup_execute_query.clone();
      41              : 
      42              :         // Pass the result_sender back to the test, together with the request params
      43              :         let params = request.into_inner();
      44              :         test_setup.send((params, result_sender)).await.unwrap();
      45              : 
      46              :         // Wire up a receiver stream to the gRPC output
      47              :         let out = ReceiverStream::new(receive);
      48              :         Ok(tonic::Response::new(
      49              :             Box::pin(out) as Self::ExecuteQueryStream
      50              :         ))
      51            2 :     }
      52              : }
      53              : 
      54            2 : pub async fn spawn_hyper_service_mock(mock: HyperServiceMock) -> (SocketAddr, oneshot::Sender<()>) {
      55            2 :     let service = HyperServiceServer::new(mock);
      56              : 
      57              :     // Bind to port `0` so the OS finds a free port
      58            2 :     let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
      59            2 :     let addr = listener.local_addr().unwrap();
      60              : 
      61            2 :     let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
      62              : 
      63            2 :     tokio::spawn(async move {
      64            2 :         tonic::transport::Server::builder()
      65            2 :             .add_service(service)
      66            2 :             .serve_with_incoming_shutdown(
      67            2 :                 tokio_stream::wrappers::TcpListenerStream::new(listener),
      68            2 :                 async { drop(shutdown_rx.await) },
      69              :             )
      70            2 :             .await
      71            0 :             .unwrap();
      72            0 :     });
      73            2 :     (addr, shutdown_tx)
      74            2 : }