140 lines
3.9 KiB
Rust
140 lines
3.9 KiB
Rust
//! Worker management integration tests
|
|
//!
|
|
//! Tests for dynamic worker add/remove operations via management API.
|
|
//! The actual worker management API uses:
|
|
//! - POST /workers - create a worker
|
|
//! - GET /workers - list workers
|
|
//! - DELETE /workers/{worker_id} - remove a worker
|
|
|
|
use axum::{
|
|
body::Body,
|
|
extract::Request,
|
|
http::{header::CONTENT_TYPE, StatusCode},
|
|
};
|
|
use serde_json::json;
|
|
use tower::ServiceExt;
|
|
|
|
use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
|
|
|
|
#[cfg(test)]
|
|
mod worker_management_tests {
|
|
use super::*;
|
|
|
|
/// Test listing workers via API
|
|
#[tokio::test]
|
|
async fn test_list_workers() {
|
|
let config = TestRouterConfig::round_robin(3900);
|
|
|
|
let ctx = AppTestContext::new_with_config(
|
|
config,
|
|
vec![
|
|
TestWorkerConfig::healthy(19900),
|
|
TestWorkerConfig::healthy(19901),
|
|
],
|
|
)
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
// List workers via GET /workers
|
|
let req = Request::builder()
|
|
.method("GET")
|
|
.uri("/workers")
|
|
.body(Body::empty())
|
|
.unwrap();
|
|
|
|
let resp = app.clone().oneshot(req).await.unwrap();
|
|
assert_eq!(
|
|
resp.status(),
|
|
StatusCode::OK,
|
|
"GET /workers should return OK"
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
|
|
/// Test that routing continues to work with multiple workers
|
|
#[tokio::test]
|
|
async fn test_routing_with_multiple_workers() {
|
|
let config = TestRouterConfig::round_robin(3901);
|
|
|
|
let ctx = AppTestContext::new_with_config(
|
|
config,
|
|
vec![
|
|
TestWorkerConfig::healthy(19902),
|
|
TestWorkerConfig::healthy(19903),
|
|
],
|
|
)
|
|
.await;
|
|
|
|
let app = ctx.create_app().await;
|
|
let mut success_count = 0;
|
|
|
|
// Verify routing distributes across workers
|
|
for i in 0..20 {
|
|
let payload = json!({
|
|
"text": format!("Test request {}", i),
|
|
"stream": false
|
|
});
|
|
|
|
let req = Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header(CONTENT_TYPE, "application/json")
|
|
.body(Body::from(serde_json::to_string(&payload).unwrap()))
|
|
.unwrap();
|
|
|
|
let resp = app.clone().oneshot(req).await.unwrap();
|
|
if resp.status() == StatusCode::OK {
|
|
success_count += 1;
|
|
}
|
|
}
|
|
|
|
assert_eq!(
|
|
success_count, 20,
|
|
"All requests should succeed with multiple workers"
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
|
|
/// Test that requests continue to work during worker operations
|
|
#[tokio::test]
|
|
async fn test_requests_during_worker_changes() {
|
|
let config = TestRouterConfig::round_robin(3902);
|
|
|
|
let ctx =
|
|
AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(19904)]).await;
|
|
|
|
let app = ctx.create_app().await;
|
|
|
|
// Send requests and verify they succeed
|
|
let mut success_count = 0;
|
|
for i in 0..10 {
|
|
let payload = json!({
|
|
"text": format!("Request during changes {}", i),
|
|
"stream": false
|
|
});
|
|
|
|
let req = Request::builder()
|
|
.method("POST")
|
|
.uri("/generate")
|
|
.header(CONTENT_TYPE, "application/json")
|
|
.body(Body::from(serde_json::to_string(&payload).unwrap()))
|
|
.unwrap();
|
|
|
|
let resp = app.clone().oneshot(req).await.unwrap();
|
|
if resp.status() == StatusCode::OK {
|
|
success_count += 1;
|
|
}
|
|
}
|
|
|
|
assert_eq!(
|
|
success_count, 10,
|
|
"All requests should succeed during normal operation"
|
|
);
|
|
|
|
ctx.shutdown().await;
|
|
}
|
|
}
|