49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
import pytest
|
|
|
|
from mlflow.demo.base import BaseDemoGenerator, DemoFeature
|
|
|
|
|
|
def test_register_and_get(fresh_registry, stub_generator):
|
|
fresh_registry.register(stub_generator)
|
|
assert fresh_registry.get(DemoFeature.TRACES) is stub_generator
|
|
|
|
|
|
def test_register_duplicate_raises(fresh_registry, stub_generator):
|
|
fresh_registry.register(stub_generator)
|
|
with pytest.raises(ValueError, match="already registered"):
|
|
fresh_registry.register(stub_generator)
|
|
|
|
|
|
def test_get_unknown_raises(fresh_registry):
|
|
with pytest.raises(ValueError, match="not found"):
|
|
fresh_registry.get("unknown")
|
|
|
|
|
|
def test_list_generators(fresh_registry, stub_generator, another_stub_generator):
|
|
assert fresh_registry.list_generators() == []
|
|
|
|
fresh_registry.register(stub_generator)
|
|
assert fresh_registry.list_generators() == [DemoFeature.TRACES]
|
|
|
|
fresh_registry.register(another_stub_generator)
|
|
assert set(fresh_registry.list_generators()) == {DemoFeature.TRACES, DemoFeature.EVALUATION}
|
|
|
|
|
|
def test_contains(fresh_registry, stub_generator):
|
|
assert DemoFeature.TRACES not in fresh_registry
|
|
|
|
fresh_registry.register(stub_generator)
|
|
assert DemoFeature.TRACES in fresh_registry
|
|
|
|
|
|
def test_register_requires_name(fresh_registry):
|
|
class NoNameGenerator(BaseDemoGenerator):
|
|
def generate(self):
|
|
pass
|
|
|
|
def _data_exists(self):
|
|
return False
|
|
|
|
with pytest.raises(ValueError, match="must define 'name'"):
|
|
fresh_registry.register(NoNameGenerator)
|