server/repository/znippy-python/tests/uv_airflow_journey.rs
//! The real Python journey, end to end, driven by a **pure-Rust client** against
//! Holger's PEP 503 / PEP 691 simple index — no shelling to `uv`/`pip`.
//!
//! Two acts, both hitting the exact `handle_http2_request` door stock `uv`/`pip`
//! reach over the wire (in-process here = no network, but byte-for-byte the same
//! protocol surface the HTTP gateway dispatches to):
//!
//! 1. **`uv`-style install** — GET the simple index, GET each project's detail
//! page, parse the wheel link + its `#sha256=` fragment, download the wheel,
//! and **independently recompute the SHA-256** to prove the advertised digest
//! matches the served bytes (content-addressed resolution, the thing pip/uv
//! abort the install on if it fails). Also cross-checks the PEP 691 JSON
//! representation carries the identical filename + hash.
//! 2. **Airflow DAG** — a 3-task DAG (`extract → transform → load`), each task
//! declaring the packages it needs; tasks run in topological order and each
//! resolves+verifies its wheels from Holger before "running". Models an
//! Airflow worker hydrating a task's Python env from the internal index.
use std::collections::BTreeMap;
use holger_znippy_python_repository::PipRepoZnippy;
use sha2::{Digest, Sha256};
use traits::RepositoryBackendTrait;
use znippy_common::plugin::PluginRegistry;
use znippy_compress::compress_dir;
use znippy_plugin_python::NativePythonPlugin;
// ── minimal real-zip wheel builder (stored entries; the python plugin parses tags) ──
fn crc32(data: &[u8]) -> u32 {
let mut crc = 0xFFFF_FFFFu32;
for &b in data {
crc ^= b as u32;
for _ in 0..8 {
crc = if crc & 1 != 0 {
(crc >> 1) ^ 0xEDB8_8320
} else {
crc >> 1
};
}
}
!crc
}
/// A one-entry STORED zip — enough to be a valid `.whl` (a zip) whose single
/// member is the given path with the given bytes.
fn build_zip_stored(name: &str, data: &[u8]) -> Vec<u8> {
let nb = name.as_bytes();
let crc = crc32(data);
let size = data.len() as u32;
let mut out = Vec::new();
let local = out.len() as u32;
out.extend_from_slice(b"PK\x03\x04");
out.extend_from_slice(&20u16.to_le_bytes());
for _ in 0..3 {
out.extend_from_slice(&0u16.to_le_bytes());
}
out.extend_from_slice(&0u16.to_le_bytes());
out.extend_from_slice(&crc.to_le_bytes());
out.extend_from_slice(&size.to_le_bytes());
out.extend_from_slice(&size.to_le_bytes());
out.extend_from_slice(&(nb.len() as u16).to_le_bytes());
out.extend_from_slice(&0u16.to_le_bytes());
out.extend_from_slice(nb);
out.extend_from_slice(data);
let cd = out.len() as u32;
out.extend_from_slice(b"PK\x01\x02");
out.extend_from_slice(&20u16.to_le_bytes());
out.extend_from_slice(&20u16.to_le_bytes());
for _ in 0..4 {
out.extend_from_slice(&0u16.to_le_bytes());
}
out.extend_from_slice(&crc.to_le_bytes());
out.extend_from_slice(&size.to_le_bytes());
out.extend_from_slice(&size.to_le_bytes());
out.extend_from_slice(&(nb.len() as u16).to_le_bytes());
for _ in 0..3 {
out.extend_from_slice(&0u16.to_le_bytes());
}
out.extend_from_slice(&0u16.to_le_bytes());
out.extend_from_slice(&0u32.to_le_bytes());
out.extend_from_slice(&local.to_le_bytes());
out.extend_from_slice(nb);
let cd_size = out.len() as u32 - cd;
out.extend_from_slice(b"PK\x05\x06");
out.extend_from_slice(&0u16.to_le_bytes());
out.extend_from_slice(&0u16.to_le_bytes());
out.extend_from_slice(&1u16.to_le_bytes());
out.extend_from_slice(&1u16.to_le_bytes());
out.extend_from_slice(&cd_size.to_le_bytes());
out.extend_from_slice(&cd.to_le_bytes());
out.extend_from_slice(&0u16.to_le_bytes());
out
}
fn sha256_hex(bytes: &[u8]) -> String {
let d = Sha256::digest(bytes);
let mut s = String::with_capacity(64);
for b in d {
s.push_str(&format!("{:02x}", b));
}
s
}
/// The demo distribution set: `(pip_name, version)`. Realistic Airflow-adjacent
/// names so the journey reads like a genuine data-science install.
fn demo_dists() -> Vec<(&'static str, &'static str)> {
vec![
("apache-airflow", "2.9.0"),
("pandas", "2.2.2"),
("numpy", "1.26.4"),
("requests", "2.32.3"),
("pendulum", "3.0.0"),
]
}
/// Wheel filename for a dist per PEP 427: dist name uses `_` for `-`.
fn wheel_filename(name: &str, version: &str) -> String {
format!("{}-{}-py3-none-any.whl", name.replace('-', "_"), version)
}
/// Build a znippy archive holding one wheel per demo dist, laid at
/// `packages/{normalized}/{filename}`, packed with the python plugin (authentic
/// path). Returns `(repo, expected_sha256_by_filename)`.
fn build_index(tmp: &tempfile::TempDir) -> (PipRepoZnippy, BTreeMap<String, String>) {
let src = tempfile::tempdir().unwrap();
let mut expected = BTreeMap::new();
for (name, version) in demo_dists() {
let norm = name.to_lowercase().replace(['_', '.'], "-");
let filename = wheel_filename(name, version);
// Distinct, non-trivial wheel bytes per dist so hashes differ.
let payload = format!(
"wheel for {name}=={version}\n{}",
"0123456789abcdef".repeat(64)
);
let wheel = build_zip_stored(
&format!("{}-{}.dist-info/METADATA", name.replace('-', "_"), version),
payload.as_bytes(),
);
expected.insert(filename.clone(), sha256_hex(&wheel));
let dir = src.path().join(format!("packages/{norm}"));
std::fs::create_dir_all(&dir).unwrap();
std::fs::write(dir.join(&filename), &wheel).unwrap();
}
let archive = tmp.path().join("index.znippy");
compress_dir(
&src.path().to_path_buf(),
&archive,
false,
Some(&PluginRegistry::with_plugin(Box::new(NativePythonPlugin))),
None,
None,
)
.expect("compress_dir with python plugin");
let repo = PipRepoZnippy::with_archive("pip".into(), archive).unwrap();
(repo, expected)
}
// ── pure-Rust index client (what uv/pip do over the protocol) ──
fn get(repo: &PipRepoZnippy, suburl: &str) -> (u16, String, Vec<u8>) {
let (status, headers, body) = repo
.handle_http2_request("GET", suburl, b"")
.expect("handle_http2_request");
let ctype = headers
.iter()
.find(|(k, _)| k.eq_ignore_ascii_case("content-type"))
.map(|(_, v)| v.clone())
.unwrap_or_default();
(status, ctype, body)
}
/// Extract `(inner_text, href)` for every `<a href="...">text</a>` in a fragment.
/// Deliberately tiny — the index HTML we emit is well-formed and flat.
fn parse_links(html: &str) -> Vec<(String, String)> {
let mut out = Vec::new();
let mut rest = html;
while let Some(a) = rest.find("<a ") {
rest = &rest[a..];
let Some(hstart) = rest.find("href=\"") else {
break;
};
let after = &rest[hstart + 6..];
let Some(hend) = after.find('"') else { break };
let href = &after[..hend];
let tail = &after[hend + 1..];
let Some(gt) = tail.find('>') else { break };
let text_area = &tail[gt + 1..];
let Some(close) = text_area.find("</a>") else {
break;
};
let text = &text_area[..close];
out.push((text.to_string(), href.to_string()));
rest = &text_area[close + 4..];
}
out
}
/// Resolve + download + verify ONE distribution the way uv/pip do: index detail →
/// pick the wheel link → split off `#sha256=` → download → recompute → compare.
/// Returns the number of bytes installed. Panics on any protocol/verification
/// failure (the test's assertion surface).
fn install_one(repo: &PipRepoZnippy, name: &str, version: &str) -> usize {
let want_file = wheel_filename(name, version);
let (status, ctype, body) = get(repo, &format!("/pip/simple/{name}/"));
assert_eq!(status, 200, "detail page for {name}");
assert!(
ctype.starts_with("text/html"),
"detail is HTML, got {ctype}"
);
let html = String::from_utf8(body).unwrap();
let links = parse_links(&html);
let (link_text, href) = links
.iter()
.find(|(text, _)| text == &want_file)
.unwrap_or_else(|| panic!("no link for {want_file} in:\n{html}"));
assert_eq!(link_text, &want_file);
// Split the PEP 503 fragment: href#sha256=<hex>.
let (path, frag) = href
.split_once("#sha256=")
.unwrap_or_else(|| panic!("link for {want_file} carries no #sha256= fragment: {href}"));
assert_eq!(frag.len(), 64, "sha256 fragment is 64 hex chars");
let (dstatus, _dctype, wheel) = get(repo, path);
assert_eq!(dstatus, 200, "download {path}");
assert!(!wheel.is_empty(), "wheel bytes non-empty");
// The verification pip/uv do: recompute and compare. Mismatch = aborted install.
let got = sha256_hex(&wheel);
assert_eq!(
&got, frag,
"content-address mismatch for {want_file}: advertised {frag}, served {got}"
);
// PEP 691 JSON cross-check: same filename + same sha256.
let (jstatus, jctype, jbody) = get(repo, &format!("/pip/simple/{name}/json"));
assert_eq!(jstatus, 200);
assert!(
jctype.contains("application/vnd.pypi.simple.v1+json"),
"PEP 691 content-type, got {jctype}"
);
let json = String::from_utf8(jbody).unwrap();
assert!(
json.contains(&format!("\"filename\":\"{want_file}\"")),
"JSON lists the wheel"
);
assert!(
json.contains(&format!("\"sha256\":\"{frag}\"")),
"JSON carries the same hash"
);
wheel.len()
}
#[test]
fn uv_style_install_resolves_and_verifies_every_wheel() {
let tmp = tempfile::tempdir().unwrap();
let (repo, expected) = build_index(&tmp);
// Act 1a: the top-level simple index lists every project (PEP 503 root).
let (status, ctype, body) = get(&repo, "/pip/simple/");
assert_eq!(status, 200);
assert!(ctype.starts_with("text/html"));
let root = String::from_utf8(body).unwrap();
for (name, _) in demo_dists() {
let norm = name.to_lowercase().replace(['_', '.'], "-");
assert!(
root.contains(&format!("{norm}/")),
"root index lists {norm}"
);
}
// Act 1b: install each dist over the protocol, verifying the content-address.
let mut total = 0usize;
for (name, version) in demo_dists() {
total += install_one(&repo, name, version);
}
assert!(total > 0, "installed wheel bytes");
// Every advertised hash matched what we recomputed (install_one asserts each),
// and the expected set is exactly the demo set.
assert_eq!(expected.len(), demo_dists().len());
}
#[test]
fn airflow_dag_hydrates_each_task_env_from_holger() {
let tmp = tempfile::tempdir().unwrap();
let (repo, _expected) = build_index(&tmp);
// A tiny DAG: (task, upstream deps, required python packages).
// Models an Airflow worker installing each task's requirements from the
// internal index at task start-up.
struct Task {
name: &'static str,
upstream: &'static [&'static str],
requires: &'static [(&'static str, &'static str)],
}
let dag = vec![
Task {
name: "extract",
upstream: &[],
requires: &[("requests", "2.32.3")],
},
Task {
name: "transform",
upstream: &["extract"],
requires: &[("pandas", "2.2.2"), ("numpy", "1.26.4")],
},
Task {
name: "load",
upstream: &["transform"],
requires: &[("apache-airflow", "2.9.0"), ("pendulum", "3.0.0")],
},
];
// Topological run: a task runs only after all upstreams have run.
let mut done: Vec<&str> = Vec::new();
let mut installed_total = 0usize;
// Simple fixpoint over the (already-ordered, but we don't rely on order) DAG.
while done.len() < dag.len() {
let before = done.len();
for t in &dag {
if done.contains(&t.name) {
continue;
}
if !t.upstream.iter().all(|u| done.contains(u)) {
continue; // upstream not finished — cannot start
}
// Task start-up: hydrate the env by pulling+verifying each requirement.
for (name, version) in t.requires {
installed_total += install_one(&repo, name, version);
}
done.push(t.name);
}
assert!(done.len() > before, "DAG made progress (no cycle/stall)");
}
assert_eq!(
done,
vec!["extract", "transform", "load"],
"ran in dependency order"
);
assert!(
installed_total > 0,
"each task hydrated its wheels from Holger"
);
}