container_registry/
test_support.rs

1//! Testing support.
2//!
3//! Requires the `test-support` feature to be enabled.
4//!
5//! This module contains utility functions to make it easier to both test the `container-registry`
6//! itself, as well as provide support when implementing tests in other crate that may need access
7//! to a container registry.
8//!
9//! ## Creating instances for testing
10//!
11//! Start by constructing a registry with the [`ContainerRegistryBuilder::build_for_testing`]
12//! method instead of the regular [`ContainerRegistryBuilder::build`] method:
13//!
14//! ```
15//! use container_registry::ContainerRegistry;
16//! use tower::util::ServiceExt;
17//!
18//! // Note: The registry is preconfigured differently when `build_for_testing` is used.
19//! let ctx = ContainerRegistry::builder().build_for_testing();
20//!
21//! // For testing of the registry itself, it can be turned into an `axum` service:/
22//! let mut service = ctx.make_service();
23//!
24//! // To launch the app and potentially use `app.call`:
25//! // let app = service.ready().await.expect("could not launch service");
26//! ```
27use std::{net::SocketAddr, sync::Arc, thread};
28
29use axum::{body::Body, routing::RouterIntoService};
30use tokio::runtime::Runtime;
31use tower_http::trace::TraceLayer;
32
33use super::{
34    auth::{self, Permissions},
35    ContainerRegistry, ContainerRegistryBuilder,
36};
37
38/// A context of a container registry instantiated for testing.
39pub struct TestingContainerRegistry {
40    /// Reference to the registry instance.
41    pub registry: Arc<ContainerRegistry>,
42    /// Storage used by the registry.
43    pub temp_storage: Option<tempdir::TempDir>,
44    /// The body limit to set when running standalone.
45    pub body_limit: usize,
46    /// The address to bind to.
47    pub bind_addr: SocketAddr,
48}
49
50/// A running registry.
51///
52/// Dropping it will cause the registry to shut down.
53pub struct RunningRegistry {
54    bound_addr: SocketAddr,
55    join_handle: Option<thread::JoinHandle<()>>,
56    _temp_storage: Option<tempdir::TempDir>,
57    shutdown: Option<tokio::sync::mpsc::Sender<()>>,
58}
59
60impl RunningRegistry {
61    /// Returns the address the registry is bound to.
62    pub fn bound_addr(&self) -> SocketAddr {
63        self.bound_addr
64    }
65}
66
67impl Drop for RunningRegistry {
68    fn drop(&mut self) {
69        // First, we signal the registry to shutdown by closing the channel:
70        drop(self.shutdown.take());
71
72        // Now we can wait for `axum` and thus the runtime and its thread to exit:
73        if let Some(join_handle) = self.join_handle.take() {
74            join_handle.join().expect("failed to join");
75        }
76
77        // All shut down, the temporary directory will be cleaned up once we exit.
78    }
79}
80
81impl TestingContainerRegistry {
82    /// Creates an `axum` service for the registry.
83    pub fn make_service(&self) -> RouterIntoService<Body> {
84        self.registry
85            .clone()
86            .make_router()
87            .layer(TraceLayer::new_for_http())
88            .into_service::<Body>()
89    }
90
91    /// Address to bind to.
92    pub fn bind(&mut self, addr: SocketAddr) -> &mut Self {
93        self.bind_addr = addr;
94        self
95    }
96
97    /// Sets the body limit, in bytes.
98    pub fn body_limit(&mut self, body_limit: usize) -> &mut Self {
99        self.body_limit = body_limit;
100        self
101    }
102
103    /// Runs a registry in a seperate thread in the background.
104    ///
105    /// Returns a handle to the registry running in the background. If dropped, the registry will
106    /// be shutdown and its storage cleaned up.
107    pub fn run_in_background(mut self) -> RunningRegistry {
108        let app = axum::Router::new()
109            .merge(self.registry.clone().make_router())
110            .layer(axum::extract::DefaultBodyLimit::max(self.body_limit));
111
112        let listener =
113            std::net::TcpListener::bind(self.bind_addr).expect("could not bind listener");
114        listener
115            .set_nonblocking(true)
116            .expect("could not set listener to non-blocking");
117        let bound_addr = listener.local_addr().expect("failed to get local address");
118
119        let (shutdown_sender, mut shutdown_receiver) = tokio::sync::mpsc::channel::<()>(1);
120        let rt = Runtime::new().expect("could not create tokio runtime");
121        let join_handle = thread::spawn(move || {
122            rt.block_on(async move {
123                let listener = tokio::net::TcpListener::from_std(listener)
124                    .expect("could not create tokio listener");
125
126                axum::serve(listener, app)
127                    .with_graceful_shutdown(async move {
128                        shutdown_receiver.recv().await;
129                    })
130                    .await
131                    .expect("axum io error");
132            })
133        });
134
135        RunningRegistry {
136            bound_addr,
137            join_handle: Some(join_handle),
138            shutdown: Some(shutdown_sender),
139            _temp_storage: self.temp_storage.take(),
140        }
141    }
142
143    /// Grants access to the registry.
144    pub fn registry(&self) -> &ContainerRegistry {
145        &self.registry
146    }
147}
148
149impl ContainerRegistryBuilder {
150    /// Constructs a new registry for testing purposes.
151    ///
152    /// Similar to [`Self::build`], except
153    ///
154    /// * If no auth provider has been set, a default one granting **full write access** to any
155    ///   user, including anonymous ones.
156    /// * If no storage path has been set, creates a temporary directory for the registry, which
157    ///   will be cleaned up if `TestingContainerRegistry` is dropped.
158    ///
159    /// # Panics
160    ///
161    /// Will panic if filesystem operations when setting up storage fail.
162    pub fn build_for_testing(mut self) -> TestingContainerRegistry {
163        let temp_storage = if self.storage.is_none() {
164            let temp_storage = tempdir::TempDir::new("container-registry-for-testing").expect(
165                "could not create temporary directory to host testing container registry instance",
166            );
167            self = self.storage(temp_storage.path());
168            Some(temp_storage)
169        } else {
170            None
171        };
172
173        if self.auth_provider.is_none() {
174            self = self.auth_provider(Arc::new(auth::Anonymous::new(
175                Permissions::ReadWrite,
176                Permissions::ReadWrite,
177            )));
178        }
179
180        let registry = self.build().expect("could not create registry");
181
182        TestingContainerRegistry {
183            registry,
184            temp_storage,
185            bind_addr: ([127, 0, 0, 1], 0).into(),
186            body_limit: 100 * 1024 * 1024,
187        }
188    }
189}