53 lines
1.5 KiB
Python
53 lines
1.5 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
|
||
|
|
|
||
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||
|
|
|
||
|
|
|
||
|
|
class FakeARDiffusionRunner:
|
||
|
|
def __init__(self) -> None:
|
||
|
|
self.resets: list[str] = []
|
||
|
|
self.closes: list[str] = []
|
||
|
|
|
||
|
|
def reset_session(self, session_id: str) -> None:
|
||
|
|
self.resets.append(session_id)
|
||
|
|
|
||
|
|
def close_session(self, session_id: str) -> None:
|
||
|
|
self.closes.append(session_id)
|
||
|
|
|
||
|
|
|
||
|
|
class FakeRegularRunner:
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
def make_worker(runner: object) -> DiffusionWorker:
|
||
|
|
worker = object.__new__(DiffusionWorker)
|
||
|
|
worker.model_runner = runner
|
||
|
|
return worker
|
||
|
|
|
||
|
|
|
||
|
|
def test_worker_delegates_ar_session_lifecycle_to_runner() -> None:
|
||
|
|
runner = FakeARDiffusionRunner()
|
||
|
|
worker = make_worker(runner)
|
||
|
|
|
||
|
|
assert worker.reset_ar_diffusion_session("world-1") is True
|
||
|
|
assert worker.close_ar_diffusion_session("world-1") is True
|
||
|
|
assert runner.resets == ["world-1"]
|
||
|
|
assert runner.closes == ["world-1"]
|
||
|
|
|
||
|
|
|
||
|
|
def test_worker_reports_unsupported_runner_without_model_assumptions() -> None:
|
||
|
|
worker = make_worker(FakeRegularRunner())
|
||
|
|
|
||
|
|
assert worker.reset_ar_diffusion_session("world-1") is False
|
||
|
|
assert worker.close_ar_diffusion_session("world-1") is False
|
||
|
|
|
||
|
|
|
||
|
|
def test_worker_rejects_empty_session_id() -> None:
|
||
|
|
worker = make_worker(FakeARDiffusionRunner())
|
||
|
|
|
||
|
|
with pytest.raises(ValueError, match="session_id"):
|
||
|
|
worker.close_ar_diffusion_session("")
|