Make the cache compatible with iterators (#1482)

Author: twiggler

.vscode/launch.json

@@ -1,4 +1,5 @@
 {
+    "version": "0.2.0",
     "configurations": [
         {
             "type": "debugpy",

dissect/target/helpers/cache.py

@@ -166,7 +166,7 @@ def call(self, *args, **kwargs) -> Any:
                         target.log.debug("", exc_info=e)
                     else:
                         target.log.info("Using cache for function: %s", self.fname)
-                        return reader
+                        return iter(reader)
                 else:
                     target.log.warning("Cache will NOT be used. File is empty: %s", cache_file)
             else:
@@ -211,7 +211,7 @@ def call(self, *args, **kwargs) -> Any:
                                 e,
                             )
                         target.log.debug("Caching to file: %s", temp_path)
-                        return CacheWriter(cache_file, temp_path, self.func(*args, **kwargs), writer)
+                        return iter(CacheWriter(cache_file, temp_path, self.func(*args, **kwargs), writer))
                     except Exception as e:
                         target.log.error("Cache will NOT be written. Failed to cache to file: %s (%s)", cache_file, e)  # noqa: TRY400
                         target.log.debug("", exc_info=e)

tests/helpers/test_cache.py

@@ -1,12 +1,17 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
+
+import pytest
 
 from dissect.target.helpers.cache import Cache
 from dissect.target.plugins.os.windows.amcache import AmcachePlugin
 from dissect.target.plugins.os.windows.ual import UalPlugin
 
 if TYPE_CHECKING:
+    from collections.abc import Callable, Iterator
+    from pathlib import Path
+
     from dissect.target.target import Target
 
 
@@ -56,3 +61,65 @@ def test_cache_filename(target_win: Target) -> None:
     assert (
         cache4.cache_path(target_win, ()).stem == "dissect.target.plugins.os.windows.ual.UalPlugin.client_access.KCk="
     )
+
+
+def test_cache_write_failure_behavior(target_bare: Target, tmp_path: Path) -> None:
+    """Specifically tests the 'Write Path' (Cache Miss) which returns a CacheWriter.
+
+    We verify that CacheWriter acts as an Iterator even when the underlying
+    plugin returns None (stops immediately).
+    """
+    target_bare._config.CACHE_DIR = str(tmp_path)
+
+    # 1. Mock Plugin with two modes
+    class MockPlugin:
+        def __init__(self, target: Target):
+            self.target = target
+
+        def success(self) -> Iterator[str]:
+            yield "success_data"
+
+        def failure(self) -> Iterator[str]:
+            if True:
+                return None
+            yield "unreachable"
+
+    plugin = MockPlugin(target_bare)
+
+    # 2. Setup Cache wrapper
+    # We force output="yield" to use LineReader/CacheWriter
+    # (mimicking the behavior of RecordWriter logic in a simpler test)
+    def create_wrapper(func: Callable[..., Iterator[str]]) -> Callable[..., Iterator[str]]:
+        cache = Cache(func)
+
+        def wrapper(*args: Any, **kwargs: Any) -> Iterator[str]:
+            return cache.call(*args, **kwargs)
+
+        wrapper.__output__ = "yield"
+        cache.wrapper = wrapper
+        return wrapper
+
+    wrap_success = create_wrapper(MockPlugin.success)
+    wrap_failure = create_wrapper(MockPlugin.failure)
+
+    # --- SCENARIO A: Success Case (Write Path) ---
+    # This creates a CacheWriter.
+    # IF CacheWriter is not wrapped in iter(), next() crashes here.
+    gen_success = wrap_success(plugin)
+
+    assert next(gen_success) == "success_data"
+    # Exhaust it to ensure file write completes
+    list(gen_success)
+
+    # --- SCENARIO B: Failure Case (Write Path) ---
+    # It creates a CacheWriter that wraps an empty generator.
+    gen_failure = wrap_failure(plugin)
+
+    # CRITICAL CHECK:
+    # 1. It must be an iterator (iter(obj) is obj)
+    # 2. calling next() should raise StopIteration, NOT TypeError
+    assert iter(gen_failure) is gen_failure
+
+    # CacheWriter should be an empty iterable
+    with pytest.raises(StopIteration):
+        next(gen_failure)