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 0 : "`{}` env variable is not set, skipping the test",
67 : ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME
68 : );
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 0 : "`{}` env variable is not set, skipping the test",
94 : ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME
95 : );
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 : 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 2 : async fn setup() -> Self {
147 2 : ensure_logging_ready();
148 2 : if env::var(ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME).is_err() {
149 2 : info!(
150 0 : "`{}` env variable is not set, skipping the test",
151 : ENABLE_REAL_AZURE_REMOTE_STORAGE_ENV_VAR_NAME
152 : );
153 2 : return Self::Disabled;
154 0 : }
155 0 :
156 0 : let max_keys_in_list_response = 10;
157 0 : let upload_tasks_count = 1 + (2 * usize::try_from(max_keys_in_list_response).unwrap());
158 :
159 0 : let enabled = EnabledAzure::setup(Some(max_keys_in_list_response)).await;
160 :
161 0 : match upload_simple_remote_data(&enabled.client, upload_tasks_count).await {
162 0 : ControlFlow::Continue(uploads) => {
163 0 : info!("Remote objects created successfully");
164 :
165 0 : Self::Enabled(AzureWithSimpleTestBlobs {
166 0 : enabled,
167 0 : remote_blobs: uploads,
168 0 : })
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 2 : }
179 :
180 2 : async fn teardown(self) {
181 2 : match self {
182 2 : Self::Disabled => {}
183 0 : Self::Enabled(ctx) | Self::UploadsFailed(_, ctx) => {
184 0 : cleanup(&ctx.enabled.client, ctx.remote_blobs).await;
185 : }
186 : }
187 2 : }
188 : }
189 :
190 0 : fn create_azure_client(
191 0 : max_keys_per_list_response: Option<i32>,
192 0 : ) -> anyhow::Result<Arc<GenericRemoteStorage>> {
193 : use rand::Rng;
194 :
195 0 : let remote_storage_azure_container = env::var("REMOTE_STORAGE_AZURE_CONTAINER").context(
196 0 : "`REMOTE_STORAGE_AZURE_CONTAINER` env var is not set, but real Azure tests are enabled",
197 0 : )?;
198 0 : let remote_storage_azure_region = env::var("REMOTE_STORAGE_AZURE_REGION").context(
199 0 : "`REMOTE_STORAGE_AZURE_REGION` env var is not set, but real Azure tests are enabled",
200 0 : )?;
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 0 : let millis = std::time::SystemTime::now()
205 0 : .duration_since(UNIX_EPOCH)
206 0 : .context("random Azure test prefix part calculation")?
207 0 : .as_millis();
208 0 :
209 0 : // because nanos can be the same for two threads so can millis, add randomness
210 0 : let random = rand::thread_rng().gen::<u32>();
211 0 :
212 0 : let remote_storage_config = RemoteStorageConfig {
213 0 : storage: RemoteStorageKind::AzureContainer(AzureConfig {
214 0 : container_name: remote_storage_azure_container,
215 0 : container_region: remote_storage_azure_region,
216 0 : prefix_in_container: Some(format!("test_{millis}_{random:08x}/")),
217 0 : concurrency_limit: NonZeroUsize::new(100).unwrap(),
218 0 : max_keys_per_list_response,
219 0 : }),
220 0 : timeout: Duration::from_secs(120),
221 0 : };
222 0 : Ok(Arc::new(
223 0 : GenericRemoteStorage::from_config(&remote_storage_config).context("remote storage init")?,
224 : ))
225 0 : }
|