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