Skip to main content

uv_test/
find_links.rs

1//! A local HTTP server that serves a directory of files as a PEP 503 flat link page.
2//!
3//! Useful for testing `--find-links` with an HTTP URL.
4
5use std::collections::HashMap;
6use std::path::Path;
7use std::sync::Arc;
8
9use wiremock::{Request, ResponseTemplate};
10
11use crate::http_server::{HttpServer, content_type_for_filename};
12use crate::vendor::{VendorArtifact, vendor_artifacts};
13
14enum FileData {
15    Bytes(Arc<[u8]>),
16    Vendor(&'static VendorArtifact),
17}
18
19impl FileData {
20    fn bytes(&self) -> anyhow::Result<Arc<[u8]>> {
21        match self {
22            Self::Bytes(bytes) => Ok(Arc::clone(bytes)),
23            Self::Vendor(artifact) => artifact.bytes(),
24        }
25    }
26}
27
28/// A running HTTP server that serves files from a directory as a flat links page.
29pub struct FindLinksServer {
30    server: HttpServer,
31}
32
33impl FindLinksServer {
34    /// Start a server that serves all files in the given directory.
35    pub fn new(directory: &Path) -> Self {
36        let mut files: HashMap<String, FileData> = HashMap::new();
37        let mut filenames: Vec<String> = Vec::new();
38
39        for entry in fs_err::read_dir(directory).expect("failed to read find-links directory") {
40            let entry = entry.expect("failed to read directory entry");
41            let path = entry.path();
42            if !path.is_file() {
43                continue;
44            }
45            let Some(filename) = path.file_name().map(|n| n.to_string_lossy().to_string()) else {
46                continue;
47            };
48            let bytes = fs_err::read(&path).expect("failed to read file");
49            files.insert(filename.clone(), FileData::Bytes(bytes.into()));
50            filenames.push(filename);
51        }
52        filenames.sort();
53
54        let files = Arc::new(files);
55        let filenames = Arc::new(filenames);
56        let server = HttpServer::start(move |request, server_uri| {
57            handle_request(request, server_uri, &files, &filenames)
58        });
59
60        Self { server }
61    }
62
63    /// Start a server that serves the pinned registry artifacts used by tests.
64    pub fn vendor() -> Self {
65        let mut files: HashMap<String, FileData> = HashMap::new();
66        let mut filenames: Vec<String> = Vec::new();
67
68        for artifact in vendor_artifacts() {
69            files.insert(artifact.filename.to_string(), FileData::Vendor(artifact));
70            filenames.push(artifact.filename.to_string());
71        }
72        filenames.sort();
73
74        let files = Arc::new(files);
75        let filenames = Arc::new(filenames);
76        let server = HttpServer::start(move |request, server_uri| {
77            handle_request(request, server_uri, &files, &filenames)
78        });
79
80        Self { server }
81    }
82
83    /// The base URL of the server (for use with `--find-links`).
84    pub fn url(&self) -> &str {
85        self.server.url()
86    }
87}
88
89fn handle_request(
90    request: &Request,
91    server_uri: &str,
92    files: &HashMap<String, FileData>,
93    filenames: &[String],
94) -> ResponseTemplate {
95    let path = request.url.path();
96
97    if path == "/" {
98        let links = filenames
99            .iter()
100            .map(|filename| format!("<a href=\"{server_uri}/{filename}\">{filename}</a>"))
101            .collect::<Vec<_>>()
102            .join("\n");
103        let html = format!("<!DOCTYPE html>\n<html><body>\n{links}\n</body></html>");
104        return ResponseTemplate::new(200).set_body_raw(html, "text/html");
105    }
106
107    let filename = path.trim_start_matches('/');
108    if let Some(file) = files.get(filename) {
109        return match file.bytes() {
110            Ok(bytes) => ResponseTemplate::new(200)
111                .set_body_raw(bytes.to_vec(), content_type_for_filename(filename)),
112            Err(error) => ResponseTemplate::new(500).set_body_string(format!("{error:#}")),
113        };
114    }
115
116    ResponseTemplate::new(404)
117}
118
119#[cfg(test)]
120mod tests {
121    use super::FindLinksServer;
122    use crate::vendor::vendor_artifacts;
123
124    #[test]
125    fn vendor_server_construction_does_not_load_artifacts() {
126        let _server = FindLinksServer::vendor();
127
128        assert!(
129            vendor_artifacts()
130                .iter()
131                .all(|artifact| !artifact.is_loaded())
132        );
133    }
134}