Line data Source code
1 : use std::collections::HashSet;
2 : use std::env;
3 : use std::num::NonZeroUsize;
4 : use std::ops::ControlFlow;
5 : use std::sync::Arc;
6 : use std::time::UNIX_EPOCH;
7 :
8 : use anyhow::Context;
9 : use remote_storage::{
10 : AzureConfig, GenericRemoteStorage, RemotePath, RemoteStorageConfig, RemoteStorageKind,
11 : };
12 : use test_context::AsyncTestContext;
13 : use tracing::info;
14 :
15 : mod common;
16 :
17 : #[path = "common/tests.rs"]
18 : mod tests_azure;
19 :
20 : use common::{cleanup, ensure_logging_ready, upload_remote_data, upload_simple_remote_data};
21 :
22 : const ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME: &str = "ENABLE_REAL_AZURE_REMOTE_STORAGE";
23 :
24 : const BASE_PREFIX: &str = "test";
25 :
26 : struct EnabledAzure {
27 : client: Arc<GenericRemoteStorage>,
28 : base_prefix: &'static str,
29 : }
30 :
31 : impl EnabledAzure {
32 0 : async fn setup(max_keys_in_list_response: Option<i32>) -> Self {
33 0 : let client = create_azure_client(max_keys_in_list_response)
34 0 : .context("Azure client creation")
35 0 : .expect("Azure client creation failed");
36 0 :
37 0 : EnabledAzure {
38 0 : client,
39 0 : base_prefix: BASE_PREFIX,
40 0 : }
41 0 : }
42 : }
43 :
44 : enum MaybeEnabledStorage {
45 : Enabled(EnabledAzure),
46 : Disabled,
47 : }
48 :
49 : #[async_trait::async_trait]
50 : impl AsyncTestContext for MaybeEnabledStorage {
51 8 : async fn setup() -> Self {
52 8 : ensure_logging_ready();
53 8 :
54 8 : if env::var(ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME).is_err() {
55 8 : info!(
56 8 : "`{}` env variable is not set, skipping the test",
57 8 : ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME
58 8 : );
59 8 : return Self::Disabled;
60 0 : }
61 0 :
62 0 : Self::Enabled(EnabledAzure::setup(None).await)
63 16 : }
64 : }
65 :
66 : enum MaybeEnabledStorageWithTestBlobs {
67 : Enabled(AzureWithTestBlobs),
68 : Disabled,
69 : UploadsFailed(anyhow::Error, AzureWithTestBlobs),
70 : }
71 :
72 : struct AzureWithTestBlobs {
73 : enabled: EnabledAzure,
74 : remote_prefixes: HashSet<RemotePath>,
75 : remote_blobs: HashSet<RemotePath>,
76 : }
77 :
78 : #[async_trait::async_trait]
79 : impl AsyncTestContext for MaybeEnabledStorageWithTestBlobs {
80 2 : async fn setup() -> Self {
81 2 : ensure_logging_ready();
82 2 : if env::var(ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME).is_err() {
83 2 : info!(
84 2 : "`{}` env variable is not set, skipping the test",
85 2 : ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME
86 2 : );
87 2 : return Self::Disabled;
88 0 : }
89 0 :
90 0 : let max_keys_in_list_response = 10;
91 0 : let upload_tasks_count = 1 + (2 * usize::try_from(max_keys_in_list_response).unwrap());
92 :
93 0 : let enabled = EnabledAzure::setup(Some(max_keys_in_list_response)).await;
94 :
95 0 : match upload_remote_data(&enabled.client, enabled.base_prefix, upload_tasks_count).await {
96 0 : ControlFlow::Continue(uploads) => {
97 0 : info!("Remote objects created successfully");
98 :
99 0 : Self::Enabled(AzureWithTestBlobs {
100 0 : enabled,
101 0 : remote_prefixes: uploads.prefixes,
102 0 : remote_blobs: uploads.blobs,
103 0 : })
104 : }
105 0 : ControlFlow::Break(uploads) => Self::UploadsFailed(
106 0 : anyhow::anyhow!("One or multiple blobs failed to upload to Azure"),
107 0 : AzureWithTestBlobs {
108 0 : enabled,
109 0 : remote_prefixes: uploads.prefixes,
110 0 : remote_blobs: uploads.blobs,
111 0 : },
112 0 : ),
113 : }
114 4 : }
115 :
116 2 : async fn teardown(self) {
117 2 : match self {
118 2 : Self::Disabled => {}
119 0 : Self::Enabled(ctx) | Self::UploadsFailed(_, ctx) => {
120 0 : cleanup(&ctx.enabled.client, ctx.remote_blobs).await;
121 : }
122 : }
123 2 : }
124 : }
125 :
126 : // NOTE: the setups for the list_prefixes test and the list_files test are very similar
127 : // However, they are not idential. The list_prefixes function is concerned with listing prefixes,
128 : // whereas the list_files function is concerned with listing files.
129 : // See `RemoteStorage::list_files` documentation for more details
130 : enum MaybeEnabledStorageWithSimpleTestBlobs {
131 : Enabled(AzureWithSimpleTestBlobs),
132 : Disabled,
133 : UploadsFailed(anyhow::Error, AzureWithSimpleTestBlobs),
134 : }
135 : struct AzureWithSimpleTestBlobs {
136 : enabled: EnabledAzure,
137 : remote_blobs: HashSet<RemotePath>,
138 : }
139 :
140 : #[async_trait::async_trait]
141 : impl AsyncTestContext for MaybeEnabledStorageWithSimpleTestBlobs {
142 2 : async fn setup() -> Self {
143 2 : ensure_logging_ready();
144 2 : if env::var(ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME).is_err() {
145 2 : info!(
146 2 : "`{}` env variable is not set, skipping the test",
147 2 : ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME
148 2 : );
149 2 : return Self::Disabled;
150 0 : }
151 0 :
152 0 : let max_keys_in_list_response = 10;
153 0 : let upload_tasks_count = 1 + (2 * usize::try_from(max_keys_in_list_response).unwrap());
154 :
155 0 : let enabled = EnabledAzure::setup(Some(max_keys_in_list_response)).await;
156 :
157 0 : match upload_simple_remote_data(&enabled.client, upload_tasks_count).await {
158 0 : ControlFlow::Continue(uploads) => {
159 0 : info!("Remote objects created successfully");
160 :
161 0 : Self::Enabled(AzureWithSimpleTestBlobs {
162 0 : enabled,
163 0 : remote_blobs: uploads,
164 0 : })
165 : }
166 0 : ControlFlow::Break(uploads) => Self::UploadsFailed(
167 0 : anyhow::anyhow!("One or multiple blobs failed to upload to Azure"),
168 0 : AzureWithSimpleTestBlobs {
169 0 : enabled,
170 0 : remote_blobs: uploads,
171 0 : },
172 0 : ),
173 : }
174 4 : }
175 :
176 2 : async fn teardown(self) {
177 2 : match self {
178 2 : Self::Disabled => {}
179 0 : Self::Enabled(ctx) | Self::UploadsFailed(_, ctx) => {
180 0 : cleanup(&ctx.enabled.client, ctx.remote_blobs).await;
181 : }
182 : }
183 2 : }
184 : }
185 :
186 0 : fn create_azure_client(
187 0 : max_keys_per_list_response: Option<i32>,
188 0 : ) -> anyhow::Result<Arc<GenericRemoteStorage>> {
189 : use rand::Rng;
190 :
191 0 : let remote_storage_azure_container = env::var("REMOTE_STORAGE_AZURE_CONTAINER").context(
192 0 : "`REMOTE_STORAGE_AZURE_CONTAINER` env var is not set, but real Azure tests are enabled",
193 0 : )?;
194 0 : let remote_storage_azure_region = env::var("REMOTE_STORAGE_AZURE_REGION").context(
195 0 : "`REMOTE_STORAGE_AZURE_REGION` env var is not set, but real Azure tests are enabled",
196 0 : )?;
197 :
198 : // due to how time works, we've had test runners use the same nanos as bucket prefixes.
199 : // millis is just a debugging aid for easier finding the prefix later.
200 0 : let millis = std::time::SystemTime::now()
201 0 : .duration_since(UNIX_EPOCH)
202 0 : .context("random Azure test prefix part calculation")?
203 0 : .as_millis();
204 0 :
205 0 : // because nanos can be the same for two threads so can millis, add randomness
206 0 : let random = rand::thread_rng().gen::<u32>();
207 0 :
208 0 : let remote_storage_config = RemoteStorageConfig {
209 0 : storage: RemoteStorageKind::AzureContainer(AzureConfig {
210 0 : container_name: remote_storage_azure_container,
211 0 : container_region: remote_storage_azure_region,
212 0 : prefix_in_container: Some(format!("test_{millis}_{random:08x}/")),
213 0 : concurrency_limit: NonZeroUsize::new(100).unwrap(),
214 0 : max_keys_per_list_response,
215 0 : }),
216 0 : };
217 0 : Ok(Arc::new(
218 0 : GenericRemoteStorage::from_config(&remote_storage_config).context("remote storage init")?,
219 : ))
220 0 : }
|
