arrow_integration_testing/flight_server_scenarios/
mod.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Collection of utilities for testing the Flight server.
19use std::net::SocketAddr;
20
21use arrow_flight::{FlightEndpoint, Location, Ticket};
22use tokio::net::TcpListener;
23
24pub mod auth_basic_proto;
25pub mod integration_test;
26pub mod middleware;
27
28type Error = Box<dyn std::error::Error + Send + Sync + 'static>;
29type Result<T = (), E = Error> = std::result::Result<T, E>;
30
31/// Listen on a port and return the address
32pub async fn listen_on(port: u16) -> Result<SocketAddr> {
33    let addr: SocketAddr = format!("0.0.0.0:{port}").parse()?;
34
35    let listener = TcpListener::bind(addr).await?;
36    let addr = listener.local_addr()?;
37
38    Ok(addr)
39}
40
41/// Create a FlightEndpoint with a ticket and location
42pub fn endpoint(ticket: &str, location_uri: impl Into<String>) -> FlightEndpoint {
43    FlightEndpoint {
44        ticket: Some(Ticket {
45            ticket: ticket.as_bytes().to_vec().into(),
46        }),
47        location: vec![Location {
48            uri: location_uri.into(),
49        }],
50        expiration_time: None,
51        app_metadata: vec![].into(),
52    }
53}