|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
3 | | -from typing import TYPE_CHECKING |
| 3 | +from typing import TYPE_CHECKING, Any |
| 4 | + |
| 5 | +import pytest |
4 | 6 |
|
5 | 7 | from dissect.target.helpers.cache import Cache |
6 | 8 | from dissect.target.plugins.os.windows.amcache import AmcachePlugin |
7 | 9 | from dissect.target.plugins.os.windows.ual import UalPlugin |
8 | 10 |
|
9 | 11 | if TYPE_CHECKING: |
| 12 | + from collections.abc import Callable, Iterator |
| 13 | + from pathlib import Path |
| 14 | + |
10 | 15 | from dissect.target.target import Target |
11 | 16 |
|
12 | 17 |
|
@@ -56,3 +61,65 @@ def test_cache_filename(target_win: Target) -> None: |
56 | 61 | assert ( |
57 | 62 | cache4.cache_path(target_win, ()).stem == "dissect.target.plugins.os.windows.ual.UalPlugin.client_access.KCk=" |
58 | 63 | ) |
| 64 | + |
| 65 | + |
| 66 | +def test_cache_write_failure_behavior(target_bare: Target, tmp_path: Path) -> None: |
| 67 | + """Specifically tests the 'Write Path' (Cache Miss) which returns a CacheWriter. |
| 68 | +
|
| 69 | + We verify that CacheWriter acts as an Iterator even when the underlying |
| 70 | + plugin returns None (stops immediately). |
| 71 | + """ |
| 72 | + target_bare._config.CACHE_DIR = str(tmp_path) |
| 73 | + |
| 74 | + # 1. Mock Plugin with two modes |
| 75 | + class MockPlugin: |
| 76 | + def __init__(self, target: Target): |
| 77 | + self.target = target |
| 78 | + |
| 79 | + def success(self) -> Iterator[str]: |
| 80 | + yield "success_data" |
| 81 | + |
| 82 | + def failure(self) -> Iterator[str]: |
| 83 | + if True: |
| 84 | + return None |
| 85 | + yield "unreachable" |
| 86 | + |
| 87 | + plugin = MockPlugin(target_bare) |
| 88 | + |
| 89 | + # 2. Setup Cache wrapper |
| 90 | + # We force output="yield" to use LineReader/CacheWriter |
| 91 | + # (mimicking the behavior of RecordWriter logic in a simpler test) |
| 92 | + def create_wrapper(func: Callable[..., Iterator[str]]) -> Callable[..., Iterator[str]]: |
| 93 | + cache = Cache(func) |
| 94 | + |
| 95 | + def wrapper(*args: Any, **kwargs: Any) -> Iterator[str]: |
| 96 | + return cache.call(*args, **kwargs) |
| 97 | + |
| 98 | + wrapper.__output__ = "yield" |
| 99 | + cache.wrapper = wrapper |
| 100 | + return wrapper |
| 101 | + |
| 102 | + wrap_success = create_wrapper(MockPlugin.success) |
| 103 | + wrap_failure = create_wrapper(MockPlugin.failure) |
| 104 | + |
| 105 | + # --- SCENARIO A: Success Case (Write Path) --- |
| 106 | + # This creates a CacheWriter. |
| 107 | + # IF CacheWriter is not wrapped in iter(), next() crashes here. |
| 108 | + gen_success = wrap_success(plugin) |
| 109 | + |
| 110 | + assert next(gen_success) == "success_data" |
| 111 | + # Exhaust it to ensure file write completes |
| 112 | + list(gen_success) |
| 113 | + |
| 114 | + # --- SCENARIO B: Failure Case (Write Path) --- |
| 115 | + # It creates a CacheWriter that wraps an empty generator. |
| 116 | + gen_failure = wrap_failure(plugin) |
| 117 | + |
| 118 | + # CRITICAL CHECK: |
| 119 | + # 1. It must be an iterator (iter(obj) is obj) |
| 120 | + # 2. calling next() should raise StopIteration, NOT TypeError |
| 121 | + assert iter(gen_failure) is gen_failure |
| 122 | + |
| 123 | + # CacheWriter should be an empty iterable |
| 124 | + with pytest.raises(StopIteration): |
| 125 | + next(gen_failure) |
0 commit comments