Skip to content

Commit 5057a63

Browse files
committed
Add context manager "postgres_manager" to enable using manager anywhere
1 parent 35c6c4c commit 5057a63

File tree

3 files changed

+56
-4
lines changed

3 files changed

+56
-4
lines changed

psqlextra/util.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
import copy
2+
3+
from contextlib import contextmanager
4+
5+
from .manager import PostgresManager
6+
7+
8+
@contextmanager
9+
def postgres_manager(model):
10+
"""Allows you to use the :see:PostgresManager with
11+
the specified model instance on the fly.
12+
13+
Arguments:
14+
model:
15+
The model or model instance to use this on.
16+
"""
17+
18+
manager = PostgresManager()
19+
manager.model = model
20+
21+
yield manager

tests/fake_model.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from psqlextra.models import PostgresModel
88

99

10-
def define_fake_model(fields=None):
10+
def define_fake_model(fields=None, model_base=PostgresModel):
1111
name = str(uuid.uuid4()).replace('-', '')[:8]
1212

1313
attributes = {
@@ -18,15 +18,15 @@ def define_fake_model(fields=None):
1818

1919
if fields:
2020
attributes.update(fields)
21-
model = type(name, (PostgresModel,), attributes)
21+
model = type(name, (model_base,), attributes)
2222

2323
return model
2424

2525

26-
def get_fake_model(fields=None):
26+
def get_fake_model(fields=None, model_base=PostgresModel):
2727
"""Creates a fake model to use during unit tests."""
2828

29-
model = define_fake_model(fields)
29+
model = define_fake_model(fields, model_base)
3030

3131
class TestProject:
3232

tests/test_manager_context.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
from django.db import models
2+
from django.test import TestCase
3+
4+
from psqlextra.util import postgres_manager
5+
6+
from .fake_model import get_fake_model
7+
8+
9+
class ManagerContextTestCase(TestCase):
10+
"""Tests whether the :see:postgres_manager context
11+
manager works properly."""
12+
13+
def test_manager_context(self):
14+
"""Tests whether the :see:postgres_manager context
15+
manager can be used to get access to :see:PostgresManager
16+
on a model that does not use it directly or inherits
17+
from :see:PostgresModel."""
18+
19+
model = get_fake_model({
20+
'myfield': models.CharField(max_length=255, unique=True)
21+
}, models.Model)
22+
23+
with postgres_manager(model) as manager:
24+
manager.upsert(
25+
conflict_target=['myfield'],
26+
fields=dict(
27+
myfield='beer'
28+
)
29+
)
30+
31+
assert manager.first().myfield == 'beer'

0 commit comments

Comments
 (0)