LCOV - code coverage report
Current view: top level - libs/remote_storage/tests/common - tests.rs (source / functions) Coverage Total Hit
Test: 32f4a56327bc9da697706839ed4836b2a00a408f.info Lines: 19.5 % 185 36
Test Date: 2024-02-07 07:37:29 Functions: 56.2 % 64 36

            Line data    Source code
       1              : use anyhow::Context;
       2              : use camino::Utf8Path;
       3              : use remote_storage::RemotePath;
       4              : use std::collections::HashSet;
       5              : use std::sync::Arc;
       6              : use test_context::test_context;
       7              : use tracing::debug;
       8              : 
       9              : use crate::common::{download_to_vec, upload_stream, wrap_stream};
      10              : 
      11              : use super::{
      12              :     MaybeEnabledStorage, MaybeEnabledStorageWithSimpleTestBlobs, MaybeEnabledStorageWithTestBlobs,
      13              : };
      14              : 
      15              : /// Tests that S3 client can list all prefixes, even if the response come paginated and requires multiple S3 queries.
      16              : /// Uses real S3 and requires [`ENABLE_REAL_S3_REMOTE_STORAGE_ENV_VAR_NAME`] and related S3 cred env vars specified.
      17              : /// See the client creation in [`create_s3_client`] for details on the required env vars.
      18              : /// If real S3 tests are disabled, the test passes, skipping any real test run: currently, there's no way to mark the test ignored in runtime with the
      19              : /// deafult test framework, see https://github.com/rust-lang/rust/issues/68007 for details.
      20              : ///
      21              : /// First, the test creates a set of S3 objects with keys `/${random_prefix_part}/${base_prefix_str}/sub_prefix_${i}/blob_${i}` in [`upload_remote_data`]
      22              : /// where
      23              : /// * `random_prefix_part` is set for the entire S3 client during the S3 client creation in [`create_s3_client`], to avoid multiple test runs interference
      24              : /// * `base_prefix_str` is a common prefix to use in the client requests: we would want to ensure that the client is able to list nested prefixes inside the bucket
      25              : ///
      26              : /// Then, verifies that the client does return correct prefixes when queried:
      27              : /// * with no prefix, it lists everything after its `${random_prefix_part}/` — that should be `${base_prefix_str}` value only
      28              : /// * with `${base_prefix_str}/` prefix, it lists every `sub_prefix_${i}`
      29              : ///
      30              : /// With the real S3 enabled and `#[cfg(test)]` Rust configuration used, the S3 client test adds a `max-keys` param to limit the response keys.
      31              : /// This way, we are able to test the pagination implicitly, by ensuring all results are returned from the remote storage and avoid uploading too many blobs to S3,
      32              : /// since current default AWS S3 pagination limit is 1000.
      33              : /// (see https://docs.aws.amazon.com/AmazonS3/latest/API/API_ListObjectsV2.html#API_ListObjectsV2_RequestSyntax)
      34              : ///
      35              : /// Lastly, the test attempts to clean up and remove all uploaded S3 files.
      36              : /// If any errors appear during the clean up, they get logged, but the test is not failed or stopped until clean up is finished.
      37            4 : #[test_context(MaybeEnabledStorageWithTestBlobs)]
      38            4 : #[tokio::test]
      39            4 : async fn pagination_should_work(ctx: &mut MaybeEnabledStorageWithTestBlobs) -> anyhow::Result<()> {
      40            4 :     let ctx = match ctx {
      41            0 :         MaybeEnabledStorageWithTestBlobs::Enabled(ctx) => ctx,
      42            4 :         MaybeEnabledStorageWithTestBlobs::Disabled => return Ok(()),
      43            0 :         MaybeEnabledStorageWithTestBlobs::UploadsFailed(e, _) => {
      44            0 :             anyhow::bail!("S3 init failed: {e:?}")
      45              :         }
      46              :     };
      47              : 
      48            0 :     let test_client = Arc::clone(&ctx.enabled.client);
      49            0 :     let expected_remote_prefixes = ctx.remote_prefixes.clone();
      50              : 
      51            0 :     let base_prefix = RemotePath::new(Utf8Path::new(ctx.enabled.base_prefix))
      52            0 :         .context("common_prefix construction")?;
      53            0 :     let root_remote_prefixes = test_client
      54            0 :         .list_prefixes(None)
      55            0 :         .await
      56            0 :         .context("client list root prefixes failure")?
      57            0 :         .into_iter()
      58            0 :         .collect::<HashSet<_>>();
      59            0 :     assert_eq!(
      60            0 :         root_remote_prefixes, HashSet::from([base_prefix.clone()]),
      61            0 :         "remote storage root prefixes list mismatches with the uploads. Returned prefixes: {root_remote_prefixes:?}"
      62              :     );
      63              : 
      64            0 :     let nested_remote_prefixes = test_client
      65            0 :         .list_prefixes(Some(&base_prefix))
      66            0 :         .await
      67            0 :         .context("client list nested prefixes failure")?
      68            0 :         .into_iter()
      69            0 :         .collect::<HashSet<_>>();
      70            0 :     let remote_only_prefixes = nested_remote_prefixes
      71            0 :         .difference(&expected_remote_prefixes)
      72            0 :         .collect::<HashSet<_>>();
      73            0 :     let missing_uploaded_prefixes = expected_remote_prefixes
      74            0 :         .difference(&nested_remote_prefixes)
      75            0 :         .collect::<HashSet<_>>();
      76            0 :     assert_eq!(
      77            0 :         remote_only_prefixes.len() + missing_uploaded_prefixes.len(), 0,
      78            0 :         "remote storage nested prefixes list mismatches with the uploads. Remote only prefixes: {remote_only_prefixes:?}, missing uploaded prefixes: {missing_uploaded_prefixes:?}",
      79              :     );
      80              : 
      81            0 :     Ok(())
      82            4 : }
      83              : 
      84              : /// Tests that S3 client can list all files in a folder, even if the response comes paginated and requirees multiple S3 queries.
      85              : /// Uses real S3 and requires [`ENABLE_REAL_S3_REMOTE_STORAGE_ENV_VAR_NAME`] and related S3 cred env vars specified. Test will skip real code and pass if env vars not set.
      86              : /// See `s3_pagination_should_work` for more information.
      87              : ///
      88              : /// First, create a set of S3 objects with keys `random_prefix/folder{j}/blob_{i}.txt` in [`upload_remote_data`]
      89              : /// Then performs the following queries:
      90              : ///    1. `list_files(None)`. This should return all files `random_prefix/folder{j}/blob_{i}.txt`
      91              : ///    2. `list_files("folder1")`.  This  should return all files `random_prefix/folder1/blob_{i}.txt`
      92            4 : #[test_context(MaybeEnabledStorageWithSimpleTestBlobs)]
      93            4 : #[tokio::test]
      94            4 : async fn list_files_works(ctx: &mut MaybeEnabledStorageWithSimpleTestBlobs) -> anyhow::Result<()> {
      95            4 :     let ctx = match ctx {
      96            0 :         MaybeEnabledStorageWithSimpleTestBlobs::Enabled(ctx) => ctx,
      97            4 :         MaybeEnabledStorageWithSimpleTestBlobs::Disabled => return Ok(()),
      98            0 :         MaybeEnabledStorageWithSimpleTestBlobs::UploadsFailed(e, _) => {
      99            0 :             anyhow::bail!("S3 init failed: {e:?}")
     100              :         }
     101              :     };
     102            0 :     let test_client = Arc::clone(&ctx.enabled.client);
     103            0 :     let base_prefix =
     104            0 :         RemotePath::new(Utf8Path::new("folder1")).context("common_prefix construction")?;
     105            0 :     let root_files = test_client
     106            0 :         .list_files(None)
     107            0 :         .await
     108            0 :         .context("client list root files failure")?
     109            0 :         .into_iter()
     110            0 :         .collect::<HashSet<_>>();
     111            0 :     assert_eq!(
     112            0 :         root_files,
     113            0 :         ctx.remote_blobs.clone(),
     114            0 :         "remote storage list_files on root mismatches with the uploads."
     115              :     );
     116            0 :     let nested_remote_files = test_client
     117            0 :         .list_files(Some(&base_prefix))
     118            0 :         .await
     119            0 :         .context("client list nested files failure")?
     120            0 :         .into_iter()
     121            0 :         .collect::<HashSet<_>>();
     122            0 :     let trim_remote_blobs: HashSet<_> = ctx
     123            0 :         .remote_blobs
     124            0 :         .iter()
     125            0 :         .map(|x| x.get_path())
     126            0 :         .filter(|x| x.starts_with("folder1"))
     127            0 :         .map(|x| RemotePath::new(x).expect("must be valid path"))
     128            0 :         .collect();
     129              :     assert_eq!(
     130              :         nested_remote_files, trim_remote_blobs,
     131            0 :         "remote storage list_files on subdirrectory mismatches with the uploads."
     132              :     );
     133            0 :     Ok(())
     134            4 : }
     135              : 
     136            4 : #[test_context(MaybeEnabledStorage)]
     137            4 : #[tokio::test]
     138            4 : async fn delete_non_exising_works(ctx: &mut MaybeEnabledStorage) -> anyhow::Result<()> {
     139            4 :     let ctx = match ctx {
     140            0 :         MaybeEnabledStorage::Enabled(ctx) => ctx,
     141            4 :         MaybeEnabledStorage::Disabled => return Ok(()),
     142              :     };
     143              : 
     144            0 :     let path = RemotePath::new(Utf8Path::new(
     145            0 :         format!("{}/for_sure_there_is_nothing_there_really", ctx.base_prefix).as_str(),
     146            0 :     ))
     147            0 :     .with_context(|| "RemotePath conversion")?;
     148              : 
     149            0 :     ctx.client.delete(&path).await.expect("should succeed");
     150            0 : 
     151            0 :     Ok(())
     152            4 : }
     153              : 
     154            4 : #[test_context(MaybeEnabledStorage)]
     155            4 : #[tokio::test]
     156            4 : async fn delete_objects_works(ctx: &mut MaybeEnabledStorage) -> anyhow::Result<()> {
     157            4 :     let ctx = match ctx {
     158            0 :         MaybeEnabledStorage::Enabled(ctx) => ctx,
     159            4 :         MaybeEnabledStorage::Disabled => return Ok(()),
     160              :     };
     161              : 
     162            0 :     let path1 = RemotePath::new(Utf8Path::new(format!("{}/path1", ctx.base_prefix).as_str()))
     163            0 :         .with_context(|| "RemotePath conversion")?;
     164              : 
     165            0 :     let path2 = RemotePath::new(Utf8Path::new(format!("{}/path2", ctx.base_prefix).as_str()))
     166            0 :         .with_context(|| "RemotePath conversion")?;
     167              : 
     168            0 :     let path3 = RemotePath::new(Utf8Path::new(format!("{}/path3", ctx.base_prefix).as_str()))
     169            0 :         .with_context(|| "RemotePath conversion")?;
     170              : 
     171            0 :     let (data, len) = upload_stream("remote blob data1".as_bytes().into());
     172            0 :     ctx.client.upload(data, len, &path1, None).await?;
     173              : 
     174            0 :     let (data, len) = upload_stream("remote blob data2".as_bytes().into());
     175            0 :     ctx.client.upload(data, len, &path2, None).await?;
     176              : 
     177            0 :     let (data, len) = upload_stream("remote blob data3".as_bytes().into());
     178            0 :     ctx.client.upload(data, len, &path3, None).await?;
     179              : 
     180            0 :     ctx.client.delete_objects(&[path1, path2]).await?;
     181              : 
     182            0 :     let prefixes = ctx.client.list_prefixes(None).await?;
     183              : 
     184            0 :     assert_eq!(prefixes.len(), 1);
     185              : 
     186            0 :     ctx.client.delete_objects(&[path3]).await?;
     187              : 
     188            0 :     Ok(())
     189            4 : }
     190              : 
     191            4 : #[test_context(MaybeEnabledStorage)]
     192            4 : #[tokio::test]
     193            4 : async fn upload_download_works(ctx: &mut MaybeEnabledStorage) -> anyhow::Result<()> {
     194            4 :     let MaybeEnabledStorage::Enabled(ctx) = ctx else {
     195            4 :         return Ok(());
     196              :     };
     197              : 
     198            0 :     let path = RemotePath::new(Utf8Path::new(format!("{}/file", ctx.base_prefix).as_str()))
     199            0 :         .with_context(|| "RemotePath conversion")?;
     200              : 
     201            0 :     let orig = bytes::Bytes::from_static("remote blob data here".as_bytes());
     202            0 : 
     203            0 :     let (data, len) = wrap_stream(orig.clone());
     204            0 : 
     205            0 :     ctx.client.upload(data, len, &path, None).await?;
     206              : 
     207              :     // Normal download request
     208            0 :     let dl = ctx.client.download(&path).await?;
     209            0 :     let buf = download_to_vec(dl).await?;
     210            0 :     assert_eq!(&buf, &orig);
     211              : 
     212              :     // Full range (end specified)
     213            0 :     let dl = ctx
     214            0 :         .client
     215            0 :         .download_byte_range(&path, 0, Some(len as u64))
     216            0 :         .await?;
     217            0 :     let buf = download_to_vec(dl).await?;
     218            0 :     assert_eq!(&buf, &orig);
     219              : 
     220              :     // partial range (end specified)
     221            0 :     let dl = ctx.client.download_byte_range(&path, 4, Some(10)).await?;
     222            0 :     let buf = download_to_vec(dl).await?;
     223            0 :     assert_eq!(&buf, &orig[4..10]);
     224              : 
     225              :     // partial range (end beyond real end)
     226            0 :     let dl = ctx
     227            0 :         .client
     228            0 :         .download_byte_range(&path, 8, Some(len as u64 * 100))
     229            0 :         .await?;
     230            0 :     let buf = download_to_vec(dl).await?;
     231            0 :     assert_eq!(&buf, &orig[8..]);
     232              : 
     233              :     // Partial range (end unspecified)
     234            0 :     let dl = ctx.client.download_byte_range(&path, 4, None).await?;
     235            0 :     let buf = download_to_vec(dl).await?;
     236            0 :     assert_eq!(&buf, &orig[4..]);
     237              : 
     238              :     // Full range (end unspecified)
     239            0 :     let dl = ctx.client.download_byte_range(&path, 0, None).await?;
     240            0 :     let buf = download_to_vec(dl).await?;
     241            0 :     assert_eq!(&buf, &orig);
     242              : 
     243            0 :     debug!("Cleanup: deleting file at path {path:?}");
     244            0 :     ctx.client
     245            0 :         .delete(&path)
     246            0 :         .await
     247            0 :         .with_context(|| format!("{path:?} removal"))?;
     248              : 
     249            0 :     Ok(())
     250            4 : }
     251              : 
     252            4 : #[test_context(MaybeEnabledStorage)]
     253            4 : #[tokio::test]
     254            4 : async fn copy_works(ctx: &mut MaybeEnabledStorage) -> anyhow::Result<()> {
     255            4 :     let MaybeEnabledStorage::Enabled(ctx) = ctx else {
     256            4 :         return Ok(());
     257              :     };
     258              : 
     259            0 :     let path = RemotePath::new(Utf8Path::new(
     260            0 :         format!("{}/file_to_copy", ctx.base_prefix).as_str(),
     261            0 :     ))
     262            0 :     .with_context(|| "RemotePath conversion")?;
     263            0 :     let path_dest = RemotePath::new(Utf8Path::new(
     264            0 :         format!("{}/file_dest", ctx.base_prefix).as_str(),
     265            0 :     ))
     266            0 :     .with_context(|| "RemotePath conversion")?;
     267              : 
     268            0 :     let orig = bytes::Bytes::from_static("remote blob data content".as_bytes());
     269            0 : 
     270            0 :     let (data, len) = wrap_stream(orig.clone());
     271            0 : 
     272            0 :     ctx.client.upload(data, len, &path, None).await?;
     273              : 
     274              :     // Normal download request
     275            0 :     ctx.client.copy_object(&path, &path_dest).await?;
     276              : 
     277            0 :     let dl = ctx.client.download(&path_dest).await?;
     278            0 :     let buf = download_to_vec(dl).await?;
     279            0 :     assert_eq!(&buf, &orig);
     280              : 
     281            0 :     debug!("Cleanup: deleting file at path {path:?}");
     282            0 :     ctx.client
     283            0 :         .delete_objects(&[path.clone(), path_dest.clone()])
     284            0 :         .await
     285            0 :         .with_context(|| format!("{path:?} removal"))?;
     286              : 
     287            0 :     Ok(())
     288            4 : }
        

Generated by: LCOV version 2.1-beta