aws_lite_rs/api/
sagemaker.rs

1//! Amazon SageMaker API client.
2//!
3//! Thin wrapper over generated ops. All URL construction and HTTP methods
4//! are in `ops::sagemaker::SagemakerOps`. This layer adds:
5//! - Ergonomic method signatures
6
7use crate::{
8    AwsHttpClient, Result,
9    ops::sagemaker::SagemakerOps,
10    types::sagemaker::{
11        ListNotebookInstancesInput, ListNotebookInstancesOutput, StopNotebookInstanceInput,
12    },
13};
14
15/// Client for the Amazon SageMaker API
16pub struct SagemakerClient<'a> {
17    ops: SagemakerOps<'a>,
18}
19
20impl<'a> SagemakerClient<'a> {
21    /// Create a new Amazon SageMaker API client
22    pub(crate) fn new(client: &'a AwsHttpClient) -> Self {
23        Self {
24            ops: SagemakerOps::new(client),
25        }
26    }
27
28    /// Returns a list of the Amazon SageMaker notebook instances in the requester's account.
29    pub async fn list_notebook_instances(
30        &self,
31        body: &ListNotebookInstancesInput,
32    ) -> Result<ListNotebookInstancesOutput> {
33        self.ops.list_notebook_instances(body).await
34    }
35
36    /// Terminates the ML compute instance. Before terminating the instance, SageMaker disconnects the ML storage volume from it.
37    pub async fn stop_notebook_instance(&self, body: &StopNotebookInstanceInput) -> Result<()> {
38        self.ops.stop_notebook_instance(body).await
39    }
40}
41
42#[cfg(test)]
43mod tests {
44    use super::*;
45    use crate::types::sagemaker::*;
46
47    #[tokio::test]
48    async fn list_notebook_instances_returns_empty() {
49        let mut mock = crate::MockClient::new();
50        mock.expect_post("/").returning_json(serde_json::json!({
51            "NotebookInstances": []
52        }));
53        let client = crate::AwsHttpClient::from_mock(mock);
54        let result = client
55            .sagemaker()
56            .list_notebook_instances(&ListNotebookInstancesInput::default())
57            .await
58            .unwrap();
59        assert_eq!(result.notebook_instances.len(), 0);
60        assert!(result.next_token.is_none());
61    }
62
63    #[tokio::test]
64    async fn list_notebook_instances_returns_instances() {
65        let mut mock = crate::MockClient::new();
66        mock.expect_post("/").returning_json(serde_json::json!({
67            "NotebookInstances": [
68                {
69                    "NotebookInstanceName": "my-notebook",
70                    "NotebookInstanceArn": "arn:aws:sagemaker:eu-central-1:123456789012:notebook-instance/my-notebook",
71                    "NotebookInstanceStatus": "InService",
72                    "InstanceType": "ml.t3.medium",
73                    "CreationTime": 1700000000.0,
74                    "LastModifiedTime": 1700001000.0,
75                    "Url": "my-notebook.notebook.eu-central-1.sagemaker.aws"
76                }
77            ],
78            "NextToken": null
79        }));
80        let client = crate::AwsHttpClient::from_mock(mock);
81        let result = client
82            .sagemaker()
83            .list_notebook_instances(&ListNotebookInstancesInput {
84                status_equals: Some("InService".to_string()),
85                ..Default::default()
86            })
87            .await
88            .unwrap();
89        assert_eq!(result.notebook_instances.len(), 1);
90        let nb = &result.notebook_instances[0];
91        assert_eq!(nb.notebook_instance_name.as_str(), "my-notebook");
92        assert_eq!(
93            nb.notebook_instance_arn.as_str(),
94            "arn:aws:sagemaker:eu-central-1:123456789012:notebook-instance/my-notebook"
95        );
96        assert_eq!(nb.notebook_instance_status.as_deref(), Some("InService"));
97        assert_eq!(nb.instance_type.as_deref(), Some("ml.t3.medium"));
98        assert_eq!(nb.creation_time, Some(1700000000.0));
99        assert_eq!(nb.last_modified_time, Some(1700001000.0));
100        assert_eq!(
101            nb.url.as_deref(),
102            Some("my-notebook.notebook.eu-central-1.sagemaker.aws")
103        );
104    }
105
106    #[tokio::test]
107    async fn stop_notebook_instance_succeeds() {
108        let mut mock = crate::MockClient::new();
109        // StopNotebookInstance returns empty body on success
110        mock.expect_post("/").returning_json(serde_json::json!({}));
111        let client = crate::AwsHttpClient::from_mock(mock);
112        let result = client
113            .sagemaker()
114            .stop_notebook_instance(&StopNotebookInstanceInput {
115                notebook_instance_name: "my-notebook".to_string(),
116            })
117            .await;
118        assert!(result.is_ok(), "Expected stop_notebook_instance to succeed");
119    }
120}
aws_lite_rs/api/sagemaker.rs

aws_lite_rs/api/
sagemaker.rs
