sqlmodel/tests/test_attribute_keyed_dict.py at 0aea9147aaa6df98b2c2f5b6cb440df778beb797 · fastapi/sqlmodel · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
import re
from enum import Enum

import pytest
from sqlalchemy.orm.collections import attribute_keyed_dict
from sqlmodel import Field, Relationship, Session, SQLModel, create_engine


def test_attribute_keyed_dict_works(clear_sqlmodel):
    class Color(str, Enum):
        Orange = "Orange"
        Blue = "Blue"

    class Child(SQLModel, table=True):
        __tablename__ = "children"

        id: int | None = Field(primary_key=True, default=None)
        parent_id: int = Field(foreign_key="parents.id")
        color: Color
        value: int

    class Parent(SQLModel, table=True):
        __tablename__ = "parents"

        id: int | None = Field(primary_key=True, default=None)
        children_by_color: dict[Color, Child] = Relationship(
            sa_relationship_kwargs={"collection_class": attribute_keyed_dict("color")}
        )

    engine = create_engine("sqlite://")
    SQLModel.metadata.create_all(engine)
    with Session(engine) as session:
        parent = Parent()
        session.add(parent)
        session.commit()
        session.refresh(parent)
        session.add(Child(parent_id=parent.id, color=Color.Orange, value=1))
        session.add(Child(parent_id=parent.id, color=Color.Blue, value=2))
        session.commit()
        session.refresh(parent)
        assert parent.children_by_color[Color.Orange].parent_id == parent.id
        assert parent.children_by_color[Color.Orange].color == Color.Orange
        assert parent.children_by_color[Color.Orange].value == 1
        assert parent.children_by_color[Color.Blue].parent_id == parent.id
        assert parent.children_by_color[Color.Blue].color == Color.Blue
        assert parent.children_by_color[Color.Blue].value == 2


def test_dict_relationship_throws_on_wrong_number_of_annotation_args(clear_sqlmodel):
    class Color(str, Enum):
        Orange = "Orange"
        Blue = "Blue"

    class Child(SQLModel, table=True):
        __tablename__ = "children"

        id: int | None = Field(primary_key=True, default=None)
        parent_id: int = Field(foreign_key="parents.id")
        color: Color
        value: int

    error_msg_fmt = "dict relationship field 'children_by_color' has {count} type arguments.  Exactly two required (e.g., dict[str, Model])"

    # One type arg
    with pytest.raises(ValueError, match=re.escape(error_msg_fmt.format(count=1))):

        class Parent(SQLModel, table=True):
            __tablename__ = "parents"

            id: int | None = Field(primary_key=True, default=None)
            children_by_color: dict[Color] = Relationship(
                sa_relationship_kwargs={
                    "collection_class": attribute_keyed_dict("color")
                }
            )

    # Three type args
    with pytest.raises(ValueError, match=re.escape(error_msg_fmt.format(count=3))):

        class Parent(SQLModel, table=True):
            __tablename__ = "parents"

            id: int | None = Field(primary_key=True, default=None)
            children_by_color: dict[Color, Child, str] = Relationship(
                sa_relationship_kwargs={
                    "collection_class": attribute_keyed_dict("color")
                }
            )