Rewrite asset processing tests to allow adding sources piecemeal.

andriyDev · andriyDev · commit a7b1ce4f6cff · 2025-11-19T20:11:05.000-08:00
diff --git a/crates/bevy_asset/src/processor/tests.rs b/crates/bevy_asset/src/processor/tests.rs
@@ -7,10 +7,8 @@ use alloc::{
     vec::Vec,
 };
 use async_lock::{RwLock, RwLockWriteGuard};
-use bevy_platform::{
-    collections::HashMap,
-    sync::{Mutex, PoisonError},
-};
+use atomicow::CowArc;
+use bevy_platform::sync::{Mutex, PoisonError};
 use bevy_reflect::TypePath;
 use core::marker::PhantomData;
 use futures_lite::AsyncWriteExt;
@@ -25,8 +23,8 @@ use bevy_tasks::BoxedFuture;
 use crate::{
     io::{
         memory::{Dir, MemoryAssetReader, MemoryAssetWriter},
-        AssetReader, AssetReaderError, AssetSourceBuilder, AssetSourceEvent, AssetSourceId,
-        AssetWatcher, PathStream, Reader,
+        AssetReader, AssetReaderError, AssetSourceBuilder, AssetSourceEvent, AssetWatcher,
+        PathStream, Reader,
     },
     processor::{
         AssetProcessor, LoadTransformAndSave, LogEntry, ProcessorState, ProcessorTransactionLog,
@@ -35,7 +33,8 @@ use crate::{
     saver::AssetSaver,
     tests::{run_app_until, CoolText, CoolTextLoader, CoolTextRon, SubText},
     transformer::{AssetTransformer, TransformedAsset},
-    Asset, AssetApp, AssetLoader, AssetMode, AssetPath, AssetPlugin, LoadContext,
+    Asset, AssetApp, AssetLoader, AssetMode, AssetPath, AssetPlugin, DefaultAssetSource,
+    LoadContext,
 };
 
 #[derive(Clone)]
@@ -49,7 +48,6 @@ struct AppWithProcessor {
     app: App,
     source_gate: Arc<RwLock<()>>,
     default_source_dirs: ProcessingDirs,
-    extra_sources_dirs: HashMap<String, ProcessingDirs>,
 }
 
 /// Similar to [`crate::io::gated::GatedReader`], but uses a lock instead of a channel to avoid
@@ -107,94 +105,83 @@ fn serialize_as_cool_text(text: &str) -> String {
     ron::ser::to_string_pretty(&cool_text_ron, PrettyConfig::new().new_line("\n")).unwrap()
 }
 
-fn create_app_with_asset_processor(extra_sources: &[String]) -> AppWithProcessor {
-    let mut app = App::new();
-    let source_gate = Arc::new(RwLock::new(()));
-
-    struct UnfinishedProcessingDirs {
-        source: Dir,
-        processed: Dir,
-        // The receiver channel for the source event sender for the unprocessed source.
-        source_event_sender_receiver:
-            async_channel::Receiver<async_channel::Sender<AssetSourceEvent>>,
-    }
+/// The processing dirs that need to be "finished" before being used.
+struct UnfinishedProcessingDirs {
+    /// The directory of the source assets.
+    source: Dir,
+    /// The directory of the processed assets.
+    processed: Dir,
+    /// The receiver channel for the source event sender for the unprocessed source.
+    source_event_sender_receiver: async_channel::Receiver<async_channel::Sender<AssetSourceEvent>>,
+}
 
-    impl UnfinishedProcessingDirs {
-        fn finish(self) -> ProcessingDirs {
-            ProcessingDirs {
-                source: self.source,
-                processed: self.processed,
-                // The processor listens for events on the source unconditionally, and we enable
-                // watching for the processed source, so both of these channels will be filled.
-                source_event_sender: self.source_event_sender_receiver.recv_blocking().unwrap(),
-            }
+impl UnfinishedProcessingDirs {
+    /// Attempts to finish the source by fetching the source event sender.
+    ///
+    /// This will block if called before adding the source.
+    fn finish(self) -> ProcessingDirs {
+        ProcessingDirs {
+            source: self.source,
+            processed: self.processed,
+            // The processor listens for events on the source unconditionally, and we enable
+            // watching for the processed source, so both of these channels will be filled.
+            source_event_sender: self.source_event_sender_receiver.recv_blocking().unwrap(),
         }
     }
+}
 
-    fn create_source(
-        app: &mut App,
-        source_id: AssetSourceId<'static>,
-        source_gate: Arc<RwLock<()>>,
-    ) -> UnfinishedProcessingDirs {
-        let source_dir = Dir::default();
-        let processed_dir = Dir::default();
+/// Creates a source, including its processed and unprocessed directories, gated on `source_gate`.
+fn create_source(source_gate: Arc<RwLock<()>>) -> (AssetSourceBuilder, UnfinishedProcessingDirs) {
+    let source_dir = Dir::default();
+    let processed_dir = Dir::default();
 
-        let source_memory_reader = LockGatedReader::new(
-            source_gate,
-            MemoryAssetReader {
-                root: source_dir.clone(),
-            },
-        );
-        let processed_memory_reader = MemoryAssetReader {
-            root: processed_dir.clone(),
-        };
-        let processed_memory_writer = MemoryAssetWriter {
-            root: processed_dir.clone(),
-        };
+    let source_memory_reader = LockGatedReader::new(
+        source_gate,
+        MemoryAssetReader {
+            root: source_dir.clone(),
+        },
+    );
+    let processed_memory_reader = MemoryAssetReader {
+        root: processed_dir.clone(),
+    };
+    let processed_memory_writer = MemoryAssetWriter {
+        root: processed_dir.clone(),
+    };
 
-        let (source_event_sender_sender, source_event_sender_receiver) = async_channel::bounded(1);
+    let (source_event_sender_sender, source_event_sender_receiver) = async_channel::bounded(1);
 
-        struct FakeWatcher;
+    struct FakeWatcher;
 
-        impl AssetWatcher for FakeWatcher {}
+    impl AssetWatcher for FakeWatcher {}
 
-        app.register_asset_source(
-            source_id,
-            AssetSourceBuilder::new(move || Box::new(source_memory_reader.clone()))
-                .with_watcher(move |sender: async_channel::Sender<AssetSourceEvent>| {
-                    source_event_sender_sender.send_blocking(sender).unwrap();
-                    Some(Box::new(FakeWatcher))
-                })
-                .with_processed_reader(move || Box::new(processed_memory_reader.clone()))
-                .with_processed_writer(move |_| Some(Box::new(processed_memory_writer.clone()))),
-        );
+    let source_builder = AssetSourceBuilder::new(move || Box::new(source_memory_reader.clone()))
+        .with_watcher(move |sender: async_channel::Sender<AssetSourceEvent>| {
+            source_event_sender_sender.send_blocking(sender).unwrap();
+            Some(Box::new(FakeWatcher))
+        })
+        .with_processed_reader(move || Box::new(processed_memory_reader.clone()))
+        .with_processed_writer(move |_| Some(Box::new(processed_memory_writer.clone())));
 
+    (
+        source_builder,
         UnfinishedProcessingDirs {
             source: source_dir,
             processed: processed_dir,
             source_event_sender_receiver,
-        }
-    }
+        },
+    )
+}
 
-    let default_source_dirs = create_source(&mut app, AssetSourceId::Default, source_gate.clone());
-
-    let extra_sources_dirs = extra_sources
-        .iter()
-        .map(|source_name| {
-            (
-                source_name.clone(),
-                create_source(
-                    &mut app,
-                    AssetSourceId::Name(source_name.clone().into()),
-                    source_gate.clone(),
-                ),
-            )
-        })
-        .collect::<Vec<_>>();
+fn create_app_with_asset_processor() -> AppWithProcessor {
+    let mut app = App::new();
+    let source_gate = Arc::new(RwLock::new(()));
+
+    let (default_source_builder, default_source_dirs) = create_source(source_gate.clone());
 
     app.add_plugins((
         TaskPoolPlugin::default(),
         AssetPlugin {
+            default_source: DefaultAssetSource::FromBuilder(Mutex::new(default_source_builder)),
             mode: AssetMode::Processed,
             use_asset_processor_override: Some(true),
             watch_for_changes_override: Some(true),
@@ -248,19 +235,25 @@ fn create_app_with_asset_processor(extra_sources: &[String]) -> AppWithProcessor
         .set_log_factory(Box::new(FakeTransactionLogFactory))
         .unwrap();
 
-    // Now that we've built the app, finish all the processing dirs.
+    // Now that we've built the app, finish the default dir, and return it.
 
     AppWithProcessor {
         app,
         source_gate,
         default_source_dirs: default_source_dirs.finish(),
-        extra_sources_dirs: extra_sources_dirs
-            .into_iter()
-            .map(|(name, dirs)| (name, dirs.finish()))
-            .collect(),
     }
 }
 
+fn register_new_source(
+    app: &mut App,
+    name: impl Into<CowArc<'static, str>>,
+    source_gate: Arc<RwLock<()>>,
+) -> ProcessingDirs {
+    let (source, processing_dirs) = create_source(source_gate);
+    app.register_asset_source(name.into(), source);
+    processing_dirs.finish()
+}
+
 fn run_app_until_finished_processing(app: &mut App, guard: RwLockWriteGuard<'_, ()>) {
     let processor = app.world().resource::<AssetProcessor>().clone();
     // We can't just wait for the processor state to be finished since we could have already
@@ -373,8 +366,7 @@ fn no_meta_or_default_processor_copies_asset() {
                 processed: processed_dir,
                 ..
             },
-        ..
-    } = create_app_with_asset_processor(&[]);
+    } = create_app_with_asset_processor();
 
     let guard = source_gate.write_blocking();
 
@@ -406,8 +398,7 @@ fn asset_processor_transforms_asset_default_processor() {
                 processed: processed_dir,
                 ..
             },
-        ..
-    } = create_app_with_asset_processor(&[]);
+    } = create_app_with_asset_processor();
 
     type CoolTextProcessor = LoadTransformAndSave<
         CoolTextLoader,
@@ -460,8 +451,7 @@ fn asset_processor_transforms_asset_with_meta() {
                 processed: processed_dir,
                 ..
             },
-        ..
-    } = create_app_with_asset_processor(&[]);
+    } = create_app_with_asset_processor();
 
     type CoolTextProcessor = LoadTransformAndSave<
         CoolTextLoader,
@@ -657,8 +647,7 @@ fn asset_processor_loading_can_read_processed_assets() {
                 processed: processed_dir,
                 ..
             },
-        ..
-    } = create_app_with_asset_processor(&[]);
+    } = create_app_with_asset_processor();
 
     // This processor loads a gltf file, converts it to BSN and then saves out the BSN.
     type GltfProcessor = LoadTransformAndSave<FakeGltfLoader, GltfToBsn, FakeBsnSaver>;
@@ -731,8 +720,7 @@ fn asset_processor_loading_can_read_source_assets() {
                 processed: processed_dir,
                 ..
             },
-        ..
-    } = create_app_with_asset_processor(&[]);
+    } = create_app_with_asset_processor();
 
     #[derive(Serialize, Deserialize)]
     struct FakeGltfxData {
@@ -927,18 +915,18 @@ fn asset_processor_processes_all_sources() {
                 processed: default_processed_dir,
                 source_event_sender: default_source_events,
             },
-        extra_sources_dirs,
-    } = create_app_with_asset_processor(&["custom_1".into(), "custom_2".into()]);
+    } = create_app_with_asset_processor();
+
     let ProcessingDirs {
         source: custom_1_source_dir,
         processed: custom_1_processed_dir,
         source_event_sender: custom_1_source_events,
-    } = extra_sources_dirs["custom_1"].clone();
+    } = register_new_source(&mut app, "custom_1", source_gate.clone());
     let ProcessingDirs {
         source: custom_2_source_dir,
         processed: custom_2_processed_dir,
         source_event_sender: custom_2_source_events,
-    } = extra_sources_dirs["custom_2"].clone();
+    } = register_new_source(&mut app, "custom_2", source_gate.clone());
 
     type AddTextProcessor = LoadTransformAndSave<
         CoolTextLoader,
@@ -1043,13 +1031,12 @@ fn nested_loads_of_processed_asset_reprocesses_on_reload() {
                 processed: default_processed_dir,
                 source_event_sender: default_source_events,
             },
-        extra_sources_dirs,
-    } = create_app_with_asset_processor(&["custom".into()]);
+    } = create_app_with_asset_processor();
     let ProcessingDirs {
         source: custom_source_dir,
         processed: custom_processed_dir,
         source_event_sender: custom_source_events,
-    } = extra_sources_dirs["custom"].clone();
+    } = register_new_source(&mut app, "custom", source_gate.clone());
 
     #[derive(Serialize, Deserialize)]
     enum NesterSerialized {
@@ -1261,8 +1248,7 @@ fn clears_invalid_data_from_processed_dir() {
                 processed: default_processed_dir,
                 ..
             },
-        ..
-    } = create_app_with_asset_processor(&[]);
+    } = create_app_with_asset_processor();
 
     type CoolTextProcessor = LoadTransformAndSave<
         CoolTextLoader,
@@ -1367,8 +1353,7 @@ fn only_reprocesses_wrong_hash_on_startup() {
             mut app,
             source_gate,
             default_source_dirs,
-            ..
-        } = create_app_with_asset_processor(&[]);
+        } = create_app_with_asset_processor();
         default_source_dir = default_source_dirs.source;
         default_processed_dir = default_source_dirs.processed;
 
@@ -1450,16 +1435,16 @@ fn only_reprocesses_wrong_hash_on_startup() {
         root: default_processed_dir.clone(),
     };
 
-    app.register_asset_source(
-        AssetSourceId::Default,
-        AssetSourceBuilder::new(move || Box::new(source_memory_reader.clone()))
-            .with_processed_reader(move || Box::new(processed_memory_reader.clone()))
-            .with_processed_writer(move |_| Some(Box::new(processed_memory_writer.clone()))),
-    );
-
     app.add_plugins((
         TaskPoolPlugin::default(),
         AssetPlugin {
+            default_source: DefaultAssetSource::FromBuilder(Mutex::new(
+                AssetSourceBuilder::new(move || Box::new(source_memory_reader.clone()))
+                    .with_processed_reader(move || Box::new(processed_memory_reader.clone()))
+                    .with_processed_writer(move |_| {
+                        Some(Box::new(processed_memory_writer.clone()))
+                    }),
+            )),
             mode: AssetMode::Processed,
             use_asset_processor_override: Some(true),
             ..Default::default()
