SQLAlchemy-Nested-Mutable/tests/test_mutable_list.py at main · wonderbeyond/SQLAlchemy-Nested-Mutable · 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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
from typing import List

import pytest
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import ARRAY, JSONB
from sqlalchemy.orm import (
    DeclarativeBase,
    Mapped,
    mapped_column,
)


from sqlalchemy_nested_mutable import MutableList, TrackedList, TrackedDict


class Base(DeclarativeBase):
    pass


class User(Base):
    __tablename__ = "user_account"

    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str] = mapped_column(sa.String(30))
    aliases = mapped_column(MutableList[str].as_mutable(ARRAY(sa.String(128))), default=list)
    schedule = mapped_column(
        MutableList[List[str]].as_mutable(ARRAY(sa.String(128), dimensions=2)), default=list
    )  # a user's weekly schedule, e.g. [ ['meeting', 'launch'], ['training', 'presentation'] ]


class UserV2(Base):
    """Use JSONB to store array data"""
    __tablename__ = "user_account_v2"

    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str] = mapped_column(sa.String(30))
    aliases = mapped_column(MutableList[str].as_mutable(JSONB), default=list)
    schedule = mapped_column(MutableList[List[str]].as_mutable(JSONB), default=list)


@pytest.fixture(scope="module", autouse=True)
def _with_tables(session):
    Base.metadata.create_all(session.bind)
    yield
    session.execute(sa.text("""
    DROP TABLE user_account CASCADE;
    DROP TABLE user_account_v2 CASCADE;
    """))
    session.commit()


def test_mutable_list(session):
    session.add(u := User(name="foo", aliases=["bar", "baz"]))
    session.commit()

    assert isinstance(u.aliases, MutableList)

    u.aliases.append("qux")
    assert u.aliases == ["bar", "baz", "qux"]

    session.commit()
    assert u.aliases == ["bar", "baz", "qux"]


def test_nested_mutable_list(session):
    session.add(u := User(
        name="foo", schedule=[["meeting", "launch"], ["training", "presentation"]]
    ))
    session.commit()

    assert isinstance(u.aliases, MutableList)
    assert u.schedule == [["meeting", "launch"], ["training", "presentation"]]

    # Mutation at top level
    u.schedule.append(["breakfast", "consulting"])
    session.commit()
    assert u.schedule == [["meeting", "launch"], ["training", "presentation"], ["breakfast", "consulting"]]

    # Mutation at nested level
    u.schedule[0][0] = "breakfast"
    session.commit()
    assert u.schedule == [["breakfast", "launch"], ["training", "presentation"], ["breakfast", "consulting"]]

    u.schedule.pop()
    session.commit()
    assert u.schedule == [["breakfast", "launch"], ["training", "presentation"]]


def test_mutable_list_stored_as_jsonb(session):
    session.add(u := UserV2(name="foo", aliases=["bar", "baz"]))
    session.commit()

    assert isinstance(u.aliases, MutableList)

    u.aliases.append("qux")
    assert u.aliases == ["bar", "baz", "qux"]

    session.commit()
    assert u.aliases == ["bar", "baz", "qux"]


def test_nested_mutable_list_stored_as_jsonb(session):
    session.add(u := UserV2(
        name="foo", schedule=[["meeting", "launch"], ["training", "presentation"]]
    ))
    session.commit()

    assert isinstance(u.aliases, MutableList)
    assert u.schedule == [["meeting", "launch"], ["training", "presentation"]]

    # Mutation at top level
    u.schedule.append(["breakfast", "consulting"])
    session.commit()
    assert u.schedule == [["meeting", "launch"], ["training", "presentation"], ["breakfast", "consulting"]]

    # Mutation at nested level
    u.schedule[0].insert(0, "breakfast")
    session.commit()
    assert u.schedule == [["breakfast", "meeting", "launch"], ["training", "presentation"], ["breakfast", "consulting"]]

    u.schedule.pop()
    session.commit()
    assert u.schedule == [["breakfast", "meeting", "launch"], ["training", "presentation"]]


def test_mutable_list_mixed_with_dict(session):
    session.add(u := UserV2(
        name="foo", schedule=[
            {"day": "mon", "events": ["meeting", "launch"]},
            {"day": "tue", "events": ["training", "presentation"]},
        ]
    ))
    session.commit()
    assert isinstance(u.schedule, MutableList)
    assert isinstance(u.schedule[0], TrackedDict)
    assert isinstance(u.schedule[0]["events"], TrackedList)

    u.schedule[0]["events"].insert(0, "breakfast")
    session.commit()
    assert u.schedule[0] == {"day": "mon", "events": ["breakfast", "meeting", "launch"]}