zeroscaler_boot_fargate/
lib.rs

1use crate::elbv2::operation::register_targets::RegisterTargetsOutput;
2use std::collections::HashMap;
3use std::env;
4use aws_config::BehaviorVersion;
5use aws_sdk_elasticloadbalancingv2 as elbv2;
6#[cfg(test)]
7use mockall::automock;
8
9const HTML_TEMPLATE: &str = r#"<!DOCTYPE html>
10<html>
11<head>
12  <meta http-equiv="refresh" content="{delay}">
13  <title>Booting Fargate...</title>
14</head>
15<body>
16  <h1>Booting {name}, please wait...</h1>
17</body>
18</html>
19"#;
20
21/// Represents the response from scaling operations
22/// 
23/// # Fields
24/// 
25/// * `status` - HTTP status code
26/// * `headers` - HTTP headers as key-value pairs
27/// * `body` - Response body content
28#[derive(Debug, PartialEq)]
29pub struct ScaleResponse {
30    pub status: u16,
31    pub headers: HashMap<String, String>,
32    pub body: String,
33}
34
35/// Implementation of the ELBv2 operations
36#[allow(dead_code)]
37pub struct Elbv2Impl {
38    inner: elbv2::Client,
39}
40
41#[cfg_attr(test, automock)]
42impl Elbv2Impl {
43    #[allow(dead_code)]
44    pub fn new(inner: elbv2::Client) -> Self {
45        Self { inner }
46    }
47
48    #[allow(dead_code)]
49    pub async fn register_target(
50        &self,
51        target_group_arn: &str,
52        target_id: &str,
53    ) -> Result<RegisterTargetsOutput, elbv2::Error> {
54        self.inner
55            .register_targets()
56            .target_group_arn(target_group_arn)
57            .targets(elbv2::types::TargetDescription::builder().id(target_id).build())
58            .send()
59            .await
60            .map_err(|e| elbv2::Error::from(e))
61    }
62}
63
64// Select the implementation based on whether we're testing
65#[cfg(test)]
66pub use MockElbv2Impl as Elbv2;
67#[cfg(not(test))]
68pub use Elbv2Impl as Elbv2;
69
70/// Registers a Fargate container with an Elastic Load Balancer target group and returns a response
71/// 
72/// # Arguments
73/// 
74/// * `target_group_arn` - The ARN of the target group to register with
75/// * `fargate_arn` - The ARN or identifier of the Fargate container to register
76/// * `delay` - The refresh delay in seconds for the HTML page
77/// * `html_template` - The HTML template to use for the response
78/// * `fargate_name` - The name of the Fargate container to display
79/// * `elbv2_client` - The ELBv2 client to use for the operation
80/// 
81/// # Returns
82/// 
83/// A `ScaleResponse` containing the HTTP response details
84///
85pub async fn scale_containers_with_params(
86    target_group_arn: &str,
87    fargate_arn: &str,
88    delay: &str,
89    html_template: &str,
90    fargate_name: &str,
91    elbv2_client: &Elbv2,
92) -> ScaleResponse {
93    if elbv2_client
94        .register_target(target_group_arn, fargate_arn)
95        .await.is_ok() {
96        let html = html_template.replace("{delay}", &delay).replace("{name}", fargate_name);
97
98        ScaleResponse {
99            status: 200,
100            headers: [("content-type".to_string(), "text/html".to_string())].into(),
101            body: html.into(),
102        }
103    } else {  
104        ScaleResponse {
105            status: 500,
106            headers: [("content-type".to_string(), "text/html".to_string())].into(),
107            body: "<!DOCTYPE html><html lang=en><meta charset=utf-8><p>Failed to register target</p>".to_string(),
108        }
109    } 
110}
111
112pub async fn scale_containers() -> ScaleResponse {
113    let target_group_arn = env::var("TARGET_GROUP_ARN")
114        .expect("TARGET_GROUP_ARN env var must be set");
115    let fargate_arn = env::var("FARGATE_ARN")
116        .expect("FARGATE_ARN env var must be set");
117    let delay = env::var("REFRESH_DELAY").unwrap_or_else(|_| "5".to_string());
118    
119    let config = aws_config::load_defaults(BehaviorVersion::latest()).await;
120    let elbv2_client = Elbv2::new(elbv2::Client::new(&config));
121    
122    scale_containers_with_params(
123        &target_group_arn,
124        &fargate_arn,
125        &delay,
126        &env::var("HTML_TEMPLATE").unwrap_or_else(|_| HTML_TEMPLATE.to_string()),
127        &env::var("FARGATE_NAME").unwrap_or_else(|_| "Fargate container".to_string()),
128        &elbv2_client,
129    ).await
130}
131
132#[cfg(test)]
133mod tests {
134    
135use super::*;
136    use std::env;
137    use aws_sdk_elasticloadbalancingv2::types::error::InvalidTargetException;
138    use mockall::predicate::*;
139
140    #[tokio::test]
141    async fn test_scale_containers_with_params() {
142        let mut mock = MockElbv2Impl::default();
143        
144        // Set up the mock expectation
145        mock.expect_register_target()
146            .with(eq("test-target-group"), eq("test-fargate"))
147            .times(1)
148            .returning(|_, _| Ok(RegisterTargetsOutput::builder().build()));
149
150        let response = scale_containers_with_params(
151            "test-target-group",
152            "test-fargate",
153            "3",
154            "<html><body>Booting {name}, please wait...</body></html>",
155            "TestFargate",
156            &mock,
157        ).await;
158
159        assert_eq!(response.status, 200);
160        assert_eq!(response.headers.get("content-type"), Some(&"text/html".to_string()));
161        assert!(response.body.contains("TestFargate"));
162    }
163
164    #[tokio::test]
165    async fn test_scale_containers_with_env_vars() {
166        let mut mock = MockElbv2Impl::default();
167        
168        mock.expect_register_target()
169            .with(eq("test-target-group"), eq("test-fargate"))
170            .times(1)
171            .returning(|_, _| Ok(RegisterTargetsOutput::builder().build()));
172unsafe {
173        env::set_var("TARGET_GROUP_ARN", "test-target-group");
174        env::set_var("FARGATE_ARN", "test-fargate");
175        env::set_var("REFRESH_DELAY", "2");
176        env::set_var("FARGATE_NAME", "TestFargate");
177}
178        // Create a wrapper function that uses our mock
179        async fn scale_containers_with_mock(mock: &Elbv2) -> ScaleResponse {
180            let target_group_arn = env::var("TARGET_GROUP_ARN").unwrap();
181            let fargate_arn = env::var("FARGATE_ARN").unwrap();
182            let delay = env::var("REFRESH_DELAY").unwrap_or_else(|_| "5".to_string());
183            
184            scale_containers_with_params(
185                &target_group_arn,
186                &fargate_arn,
187                &delay,
188                &env::var("HTML_TEMPLATE").unwrap_or_else(|_| HTML_TEMPLATE.to_string()),
189                &env::var("FARGATE_NAME").unwrap_or_else(|_| "Fargate container".to_string()),
190                mock,
191            ).await
192        }
193
194        let response = scale_containers_with_mock(&mock).await;
195
196        assert_eq!(response.status, 200);
197        assert_eq!(response.headers.get("content-type"), Some(&"text/html".to_string()));
198        assert!(response.body.contains("TestFargate"));
199    }
200
201    #[tokio::test]
202    async fn test_scale_containers_with_defaults() {
203        let mut mock = MockElbv2Impl::default();
204        
205        mock.expect_register_target()
206            .with(eq("test-target-group"), eq("test-fargate"))
207            .times(1)
208            .returning(|_, _| Ok(RegisterTargetsOutput::builder().build()));
209unsafe {
210    env::set_var("TARGET_GROUP_ARN", "test-target-group");
211    env::set_var("FARGATE_ARN", "test-fargate");
212}
213        async fn scale_containers_with_mock(mock: &Elbv2) -> ScaleResponse {
214            let target_group_arn = env::var("TARGET_GROUP_ARN").unwrap();
215            let fargate_arn = env::var("FARGATE_ARN").unwrap();
216            let delay = env::var("REFRESH_DELAY").unwrap_or_else(|_| "5".to_string());
217            
218            scale_containers_with_params(
219                &target_group_arn,
220                &fargate_arn,
221                &delay,
222                &env::var("HTML_TEMPLATE").unwrap_or_else(|_| HTML_TEMPLATE.to_string()),
223                &env::var("FARGATE_NAME").unwrap_or_else(|_| "Fargate container".to_string()),
224                mock,
225            ).await
226        }
227
228        let response = scale_containers_with_mock(&mock).await;
229
230        assert_eq!(response.status, 200);
231        assert_eq!(response.headers.get("content-type"), Some(&"text/html".to_string()));
232        assert!(response.body.contains("<h1>Booting TestFargate, please wait...</h1>"));
233    }
234
235    #[tokio::test]
236    async fn test_register_target_failure() {
237        let mut mock = MockElbv2Impl::default();
238        
239        mock.expect_register_target()
240            .with(eq("test-target-group"), eq("test-fargate"))
241            .times(1)
242            .returning(|_, _| Err(elbv2::Error::InvalidTargetException(InvalidTargetException::builder().build())
243            ));
244
245        let result = scale_containers_with_params(
246            "test-target-group",
247            "test-fargate",
248            "3",
249            "<html><body>Booting {name}, please wait...</body></html>",
250            "TestFargate",
251            &mock,
252        ).await;
253
254        // The function should panic with the expect message
255        assert_eq!(result.status , 500);
256        assert_eq!(result.headers.get("content-type"), Some(&"text/html".to_string()));
257        assert!(result.body.contains("Failed to register target"));
258    }
259}
zeroscaler_boot_fargate/lib.rs

zeroscaler_boot_fargate/
lib.rs
