def test_diagram_comparisons():
diagram1 = erd.create(Party)
diagram2 = erd.create(Party)
diagram3 = erd.create(Quest)
assert diagram1 == diagram2
assert diagram1 != diagram3
assert diagram1 in [diagram2]
assert diagram1 in {diagram2}
assert diagram1 not in [diagram3]
def test_unknown_model_type_error():
class BadModel:
pass
with pytest.raises(UnknownModelTypeError) as e_info:
erd.create(BadModel)
e = e_info.value
assert e.model is BadModel
assert "Given model does not match any supported types" in str(e)
assert BadModel.__name__ in str(e)
def create_assets(examples: ModuleType):
stem = examples.__name__.rsplit(".", 1)[1]
(ASSETS_DIR / stem).mkdir(exist_ok=True)
diagram = erd.create(examples.Party)
diagram.draw(out=ASSETS_DIR / stem / "diagram.png")
diagram.draw(out=ASSETS_DIR / stem / "diagram.svg")
with (ASSETS_DIR / stem / "diagram.dot").open("w") as fp:
fp.write(diagram.to_dot())
def test_draw_with_terminus(tmp_path):
# use EntityRelationshipDiagram.draw as expected
expected_path = tmp_path / "expected.png"
diagram = erd.create(Party, termini=[Quest])
diagram.draw(expected_path)
path = tmp_path / "diagram.png"
erd.draw(Party, out=path, termini=[Quest])
assert imghdr.what(path) == "png"
assert filecmp.cmp(path, expected_path)
def test_edge_comparisons():
diagram = erd.create(Party)
edges = diagram.edges
edge0_copy = Edge(source=edges[0].source,
source_field=edges[0].source_field,
target=edges[0].target)
assert edges[0] == edge0_copy
assert edges[1] != edge0_copy
assert edges[0] in [edge0_copy]
assert edges[0] in {edge0_copy}
assert edges[0] not in {edges[1]}
def test_string_forward_ref():
@dataclasses.dataclass
class WithStringForwardRef:
sibling: "WithStringForwardRef"
get_type_hints(WithStringForwardRef, localns=locals())
# Unevaluated forward ref should error
with pytest.raises(StringForwardRefError):
_ = erd.create(WithStringForwardRef)
@dataclasses.dataclass
class WithExplicitForwardRef:
sibling: ForwardRef("WithExplicitForwardRef") # noqa: F821
get_type_hints(WithExplicitForwardRef, localns=locals())
diagram = erd.create(WithExplicitForwardRef)
assert {m.model for m in diagram.models} == {WithExplicitForwardRef}
assert {(e.source.model, e.target.model)
for e in diagram.edges} == {(WithExplicitForwardRef,
WithExplicitForwardRef)}
def test_unevaluated_forward_ref():
class PydanticItem(BaseModel):
name: str
class PydanticContainer(BaseModel):
items: List["PydanticItem"]
# Unevaluated forward ref should error
with pytest.raises(UnevaluatedForwardRefError,
match="update_forward_refs"):
_ = erd.create(PydanticContainer)
# Evaluate forward ref
PydanticContainer.update_forward_refs(**locals())
# Test that model can be used
_ = PydanticContainer(items=[PydanticItem(name="thingie")])
diagram = erd.create(PydanticContainer)
assert {m.model
for m in diagram.models} == {PydanticContainer, PydanticItem}
assert {(e.source.model, e.target.model)
for e in diagram.edges} == {(PydanticContainer, PydanticItem)}
def test_edge_comparisons():
diagram = erd.create(Party)
edges = diagram.edges
edge0_copy = Edge(source=edges[0].source,
source_field=edges[0].source_field,
target=edges[0].target)
assert edges[0] == edge0_copy
assert edges[1] != edge0_copy
assert edges[0] in [edge0_copy]
assert edges[0] in {edge0_copy}
assert edges[0] not in {edges[1]}
# Bad comparison should raise error
with pytest.raises(TypeError, match="not supported"):
edges[0] < str(edges[0])
def test_unevaluated_forward_ref():
@dataclasses.dataclass
class DataClassItem:
name: str
@dataclasses.dataclass
class DataClassContainer:
items: List["DataClassItem"]
# Unevaluated forward ref should error
with pytest.raises(UnevaluatedForwardRefError, match="get_type_hints"):
_ = erd.create(DataClassContainer)
# Evaluate forward ref
get_type_hints(DataClassContainer, localns=locals())
# Test that class can be initialized
_ = DataClassContainer(items=[DataClassItem(name="thingie")])
diagram = erd.create(DataClassContainer)
assert {m.model
for m in diagram.models} == {DataClassContainer, DataClassItem}
assert {(e.source.model, e.target.model)
for e in diagram.edges} == {(DataClassContainer, DataClassItem)}
def test_model_graph_search_nested_args():
class Inner0(BaseModel):
id: int
class Inner1(BaseModel):
id: int
class Outer(BaseModel):
inner: Dict[str, Tuple[Inner0, Inner1]]
diagram = erd.create(Outer)
assert {m.model for m in diagram.models} == {Outer, Inner0, Inner1}
assert {(e.source.model, e.target.model)
for e in diagram.edges} == {
(Outer, Inner0),
(Outer, Inner1),
}
def test_repr_svg():
diagram = erd.create(Party)
svg = diagram._repr_svg_()
assert svg and isinstance(svg, str)
def test_repr_png():
diagram = erd.create(Party)
png = diagram._repr_png_()
assert png and isinstance(png, bytes)
assert imghdr.what("", h=png) == "png"
def test_repr():
diagram = erd.create(Party)
assert repr(diagram) and isinstance(repr(diagram), str)
assert repr(diagram.edges[0]) and isinstance(repr(diagram.edges[0]), str)
def test_to_dot_with_terminus(tmp_path):
# use EntityRelationshipDiagram.to_dot as expected
diagram = erd.create(Party, termini=[Quest])
assert erd.to_dot(Party, termini=[Quest]) == diagram.to_dot()
def test_not_a_type_error():
with pytest.raises(NotATypeError, match="Given model is not a type"):
erd.create(5)
def test_repr():
diagram = erd.create(Party)
assert repr(diagram.models[0]) and isinstance(repr(diagram.models[0]), str)
assert repr(diagram.models[0].fields[0]) and isinstance(repr(diagram.models[0].fields[0]), str)
