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 6 : async fn setup(max_keys_in_list_response: Option<i32>) -> Self {
33 6 : let client = create_azure_client(max_keys_in_list_response)
34 6 : .context("Azure client creation")
35 6 : .expect("Azure client creation failed");
36 6 :
37 6 : EnabledAzure {
38 6 : client,
39 6 : base_prefix: BASE_PREFIX,
40 6 : }
41 6 : }
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 16 : async fn setup() -> Self {
62 16 : ensure_logging_ready();
63 16 :
64 16 : if env::var(ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME).is_err() {
65 12 : info!(
66 0 : "`{}` env variable is not set, skipping the test",
67 : ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME
68 : );
69 12 : return Self::Disabled;
70 4 : }
71 4 :
72 4 : Self::Enabled(EnabledAzure::setup(None).await)
73 16 : }
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 4 : async fn setup() -> Self {
90 4 : ensure_logging_ready();
91 4 : if env::var(ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME).is_err() {
92 3 : info!(
93 0 : "`{}` env variable is not set, skipping the test",
94 : ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME
95 : );
96 3 : return Self::Disabled;
97 1 : }
98 1 :
99 1 : let max_keys_in_list_response = 10;
100 1 : let upload_tasks_count = 1 + (2 * usize::try_from(max_keys_in_list_response).unwrap());
101 :
102 1 : let enabled = EnabledAzure::setup(Some(max_keys_in_list_response)).await;
103 :
104 21 : match upload_remote_data(&enabled.client, enabled.base_prefix, upload_tasks_count).await {
105 1 : ControlFlow::Continue(uploads) => {
106 1 : info!("Remote objects created successfully");
107 :
108 1 : Self::Enabled(AzureWithTestBlobs {
109 1 : enabled,
110 1 : remote_prefixes: uploads.prefixes,
111 1 : remote_blobs: uploads.blobs,
112 1 : })
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 4 : }
124 :
125 4 : async fn teardown(self) {
126 4 : match self {
127 3 : Self::Disabled => {}
128 1 : Self::Enabled(ctx) | Self::UploadsFailed(_, ctx) => {
129 21 : cleanup(&ctx.enabled.client, ctx.remote_blobs).await;
130 : }
131 : }
132 4 : }
133 : }
134 :
135 : enum MaybeEnabledStorageWithSimpleTestBlobs {
136 : Enabled(AzureWithSimpleTestBlobs),
137 : Disabled,
138 : UploadsFailed(anyhow::Error, AzureWithSimpleTestBlobs),
139 : }
140 : struct AzureWithSimpleTestBlobs {
141 : enabled: EnabledAzure,
142 : remote_blobs: HashSet<RemotePath>,
143 : }
144 :
145 : impl AsyncTestContext for MaybeEnabledStorageWithSimpleTestBlobs {
146 4 : async fn setup() -> Self {
147 4 : ensure_logging_ready();
148 4 : if env::var(ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME).is_err() {
149 3 : info!(
150 0 : "`{}` env variable is not set, skipping the test",
151 : ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME
152 : );
153 3 : return Self::Disabled;
154 1 : }
155 1 :
156 1 : let max_keys_in_list_response = 10;
157 1 : let upload_tasks_count = 1 + (2 * usize::try_from(max_keys_in_list_response).unwrap());
158 :
159 1 : let enabled = EnabledAzure::setup(Some(max_keys_in_list_response)).await;
160 :
161 21 : match upload_simple_remote_data(&enabled.client, upload_tasks_count).await {
162 1 : ControlFlow::Continue(uploads) => {
163 1 : info!("Remote objects created successfully");
164 :
165 1 : Self::Enabled(AzureWithSimpleTestBlobs {
166 1 : enabled,
167 1 : remote_blobs: uploads,
168 1 : })
169 : }
170 0 : ControlFlow::Break(uploads) => Self::UploadsFailed(
171 0 : anyhow::anyhow!("One or multiple blobs failed to upload to Azure"),
172 0 : AzureWithSimpleTestBlobs {
173 0 : enabled,
174 0 : remote_blobs: uploads,
175 0 : },
176 0 : ),
177 : }
178 4 : }
179 :
180 4 : async fn teardown(self) {
181 4 : match self {
182 3 : Self::Disabled => {}
183 1 : Self::Enabled(ctx) | Self::UploadsFailed(_, ctx) => {
184 21 : cleanup(&ctx.enabled.client, ctx.remote_blobs).await;
185 : }
186 : }
187 4 : }
188 : }
189 :
190 6 : fn create_azure_client(
191 6 : max_keys_per_list_response: Option<i32>,
192 6 : ) -> anyhow::Result<Arc<GenericRemoteStorage>> {
193 : use rand::Rng;
194 :
195 6 : let remote_storage_azure_container = env::var("REMOTE_STORAGE_AZURE_CONTAINER").context(
196 6 : "`REMOTE_STORAGE_AZURE_CONTAINER` env var is not set, but real Azure tests are enabled",
197 6 : )?;
198 6 : let remote_storage_azure_region = env::var("REMOTE_STORAGE_AZURE_REGION").context(
199 6 : "`REMOTE_STORAGE_AZURE_REGION` env var is not set, but real Azure tests are enabled",
200 6 : )?;
201 :
202 : // due to how time works, we've had test runners use the same nanos as bucket prefixes.
203 : // millis is just a debugging aid for easier finding the prefix later.
204 6 : let millis = std::time::SystemTime::now()
205 6 : .duration_since(UNIX_EPOCH)
206 6 : .context("random Azure test prefix part calculation")?
207 6 : .as_millis();
208 6 :
209 6 : // because nanos can be the same for two threads so can millis, add randomness
210 6 : let random = rand::thread_rng().gen::<u32>();
211 6 :
212 6 : let remote_storage_config = RemoteStorageConfig {
213 6 : storage: RemoteStorageKind::AzureContainer(AzureConfig {
214 6 : container_name: remote_storage_azure_container,
215 6 : storage_account: None,
216 6 : container_region: remote_storage_azure_region,
217 6 : prefix_in_container: Some(format!("test_{millis}_{random:08x}/")),
218 6 : concurrency_limit: NonZeroUsize::new(100).unwrap(),
219 6 : max_keys_per_list_response,
220 6 : }),
221 6 : timeout: Duration::from_secs(120),
222 6 : };
223 6 : Ok(Arc::new(
224 6 : GenericRemoteStorage::from_config(&remote_storage_config).context("remote storage init")?,
225 : ))
226 6 : }
|
