aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--generalresearch/gr_django/migrations/0004_user_avatar_url_user_name_alter_user_sub.py28
-rw-r--r--generalresearch/gr_django/models.py8
-rw-r--r--generalresearch/managers/gr/authentication.py142
-rw-r--r--generalresearch/models/gr/authentication.py103
-rw-r--r--test_utils/models/gr/conftest.py25
-rw-r--r--tests/managers/gr/test_authentication.py75
-rw-r--r--tests/managers/gr/test_team.py43
7 files changed, 135 insertions, 289 deletions
diff --git a/generalresearch/gr_django/migrations/0004_user_avatar_url_user_name_alter_user_sub.py b/generalresearch/gr_django/migrations/0004_user_avatar_url_user_name_alter_user_sub.py
new file mode 100644
index 0000000..586b5f6
--- /dev/null
+++ b/generalresearch/gr_django/migrations/0004_user_avatar_url_user_name_alter_user_sub.py
@@ -0,0 +1,28 @@
+# Generated by Django 6.1.1 on 2026-10-27 18:01
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+ dependencies = [
+ ('common', '0003_user_email_alter_membership_owner_and_more'),
+ ]
+
+ operations = [
+ migrations.AddField(
+ model_name='user',
+ name='avatar_url',
+ field=models.URLField(null=True),
+ ),
+ migrations.AddField(
+ model_name='user',
+ name='name',
+ field=models.CharField(max_length=200, null=True),
+ ),
+ migrations.AlterField(
+ model_name='user',
+ name='sub',
+ field=models.CharField(max_length=200, null=True),
+ ),
+ ]
diff --git a/generalresearch/gr_django/models.py b/generalresearch/gr_django/models.py
index 7f26eaf..7d2b7b1 100644
--- a/generalresearch/gr_django/models.py
+++ b/generalresearch/gr_django/models.py
@@ -5,8 +5,10 @@ from django.db import models
class User(models.Model):
id = models.BigAutoField(primary_key=True)
- sub = models.CharField(max_length=200, blank=False, null=False)
+ sub = models.CharField(max_length=200, blank=False, null=True)
email = models.EmailField(null=True, unique=True)
+ name = models.CharField(max_length=200, null=True)
+ avatar_url = models.URLField(null=True)
is_superuser = models.BooleanField(default=False)
date_joined = models.DateTimeField(auto_now_add=True)
@@ -50,11 +52,11 @@ class Membership(models.Model):
class Meta:
constraints = [
models.UniqueConstraint(
- fields=['user_id', 'team_id'],
- name='uniq_membership_user_team'
+ fields=["user_id", "team_id"], name="uniq_membership_user_team"
)
]
+
class Team(models.Model):
"""
This is a collection of users / accounts that work as a team and share
diff --git a/generalresearch/managers/gr/authentication.py b/generalresearch/managers/gr/authentication.py
index 18c8ba4..9263f65 100644
--- a/generalresearch/managers/gr/authentication.py
+++ b/generalresearch/managers/gr/authentication.py
@@ -4,29 +4,27 @@ import binascii
import logging
import os
from datetime import UTC, datetime
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING
from psycopg import sql
-from pydantic import AnyHttpUrl, PositiveInt
+from pydantic import PositiveInt
from generalresearch.managers.base import PostgresManager, PostgresManagerWithRedis
from generalresearch.models.custom_types import UUIDStr
if TYPE_CHECKING:
-
from generalresearch.models.gr.authentication import GRToken, GRUser
from generalresearch.pg_helper import PostgresConfig
- from generalresearch.redis_helper import RedisConfig
LOG = logging.getLogger("gr")
class GRUserManager(PostgresManagerWithRedis):
-
def create(
self,
- sub: str,
- email: str | None = None,
+ email: str,
+ name: str,
+ avatar_url: str | None = None,
is_superuser: bool = False,
) -> GRUser:
from generalresearch.models.gr.authentication import GRUser
@@ -35,8 +33,9 @@ class GRUserManager(PostgresManagerWithRedis):
instance = GRUser.model_validate(
{
- "sub": sub,
"email": email,
+ "name": name,
+ "avatar_url": avatar_url,
"is_superuser": is_superuser,
"date_joined": now,
}
@@ -46,10 +45,12 @@ class GRUserManager(PostgresManagerWithRedis):
with self.pg_config.make_connection() as conn:
with conn.cursor() as c:
query = sql.SQL("""
- INSERT INTO gr_user
- (sub, email, is_superuser, date_joined)
- VALUES (%(sub)s, %(email)s, %(is_superuser)s, %(date_joined)s)
- RETURNING id
+ INSERT INTO gr_user (
+ email, name, is_superuser, date_joined, avatar_url
+ ) VALUES (
+ %(email)s, %(name)s, %(is_superuser)s,
+ %(date_joined)s, %(avatar_url)s
+ ) RETURNING id;
""")
c.execute(query=query, params=data)
gr_user_id: int = c.fetchone()["id"]
@@ -62,7 +63,6 @@ class GRUserManager(PostgresManagerWithRedis):
self,
*,
gr_user_id: int | None = None,
- sub: str | None = None,
email: str | None = None,
) -> list[GRUser]:
from generalresearch.models.gr.authentication import GRUser
@@ -72,19 +72,16 @@ class GRUserManager(PostgresManagerWithRedis):
if gr_user_id is not None:
conditions.append("u.id = %(gr_user_id)s")
- params['gr_user_id'] = gr_user_id
- if sub is not None:
- conditions.append("u.sub = %(sub)s")
- params['sub'] = sub
+ params["gr_user_id"] = gr_user_id
if email is not None:
conditions.append("u.email = %(email)s")
- params['email'] = email
+ params["email"] = email
filter_str = " AND ".join(conditions)
filter_str = "WHERE " + filter_str if filter_str else ""
query = f"""
- SELECT u.*
+ SELECT u.id, u.name, u.email, u.is_superuser, u.date_joined, u.avatar_url
FROM gr_user AS u
{filter_str}
"""
@@ -108,26 +105,11 @@ class GRUserManager(PostgresManagerWithRedis):
def get_by_id(self, gr_user_id: int) -> GRUser | None:
return self._get_one(gr_user_id=gr_user_id)
- def get_by_sub(self, sub: str, raises=True) -> GRUser | None:
- return self._get_one(sub=sub, raises=raises)
-
def get_by_email(self, email: str, raises=True) -> GRUser | None:
return self._get_one(email=email, raises=raises)
- def get_by_sub_or_create(self, sub: str) -> GRUser:
- return self.get_by_sub(sub=sub, raises=False) or self.create(sub=sub)
-
def get_all(self) -> list[GRUser]:
- from generalresearch.models.gr.authentication import GRUser
-
- with self.pg_config.make_connection() as conn, conn.cursor() as c:
- c.execute(query="""
- SELECT u.*
- FROM gr_user AS u
- """)
- res = c.fetchall()
-
- return [GRUser.from_postgresql(i) for i in res]
+ return self.filter()
def get_by_team(self, team_id: PositiveInt) -> list[GRUser]:
from generalresearch.models.gr.authentication import GRUser
@@ -135,50 +117,24 @@ class GRUserManager(PostgresManagerWithRedis):
with self.pg_config.make_connection() as conn, conn.cursor() as c:
c.execute(
query="""
- SELECT gru.*
+ SELECT u.id, u.name, u.email, u.avatar_url,
+ u.is_superuser, u.date_joined
FROM common_membership AS membership
- INNER JOIN gr_user AS gru
- ON gru.id = membership.user_id
+ INNER JOIN gr_user AS u
+ ON u.id = membership.user_id
WHERE membership.team_id = %s
""",
params=(team_id,),
)
res = c.fetchall()
- for item in res:
- for k, v in item.items():
- if isinstance(v, datetime):
- item[k] = item[k].replace(tzinfo=UTC)
-
return [GRUser.model_validate(item) for item in res]
- def list_product_uuids(
- self, user: GRUser, thl_pg_config: PostgresConfig
- ) -> list[UUIDStr] | None:
- if user.business_uuids is None:
- LOG.warning("prefetch not run")
- return None
-
- res = thl_pg_config.execute_sql_query(
- query="""
- SELECT bp.id::uuid as uuid
- FROM userprofile_brokerageproduct AS bp
- WHERE bp.business_id = ANY(%s)
- """,
- params=[user.business_uuids],
- )
- return [item["uuid"] for item in res]
-
class GRTokenManager(PostgresManager):
-
def get_by_key(
self,
api_key: str,
- jwks: dict[str, Any] | None = None,
- audience: str | None = None,
- issuer: AnyHttpUrl | str | None = None,
- gr_redis_config: RedisConfig | None = None,
) -> GRToken:
"""Return the GRToken for this API Token.
@@ -191,46 +147,16 @@ class GRTokenManager(PostgresManager):
:return GRToken instance (minified version, no relationships)
:raises NotFoundException
"""
- from generalresearch.models.gr.authentication import Claims, GRToken
-
- # SSO Key
- if GRToken.is_sso(api_key):
- from jose import jwt
-
- payload = jwt.decode(
- token=api_key,
- key=jwks,
- algorithms=["RS256"],
- audience=audience,
- issuer=issuer,
- )
- claims = Claims.model_validate(payload)
-
- gr_um = GRUserManager(
- pg_config=self.pg_config, redis_config=gr_redis_config
- )
- gr_user = gr_um.get_by_sub_or_create(sub=claims.subject)
- gr_user.claims = claims
-
- gr_token = GRToken.model_validate(
- {
- "key": api_key,
- "user_id": gr_user.id,
- "user": gr_user,
- "created": datetime.now(tz=UTC),
- }
- )
-
- return gr_token
+ from generalresearch.models.gr.authentication import GRToken
# API Key
with self.pg_config.make_connection() as conn, conn.cursor() as c:
query = sql.SQL("""
- SELECT grk.*
- FROM gr_token AS grk
- WHERE grk.key = %s
- LIMIT 1
- """)
+ SELECT grk.*
+ FROM gr_token AS grk
+ WHERE grk.key = %s
+ LIMIT 2
+ """)
c.execute(query=query, params=(api_key,))
res = c.fetchall()
@@ -240,9 +166,7 @@ class GRTokenManager(PostgresManager):
if len(res) > 1:
raise ValueError(f"Too many GRUsers found with token of '{api_key}'")
- item = res[0]
-
- return GRToken.model_validate(item)
+ return GRToken.model_validate(res[0])
def create(self, user_id: PositiveInt) -> None:
# Taken directly from the DRF Token
@@ -285,18 +209,10 @@ class GRTokenManager(PostgresManager):
WHERE u.id = %s
LIMIT 1;
""")
-
c.execute(query=query, params=(user_id,))
-
result = c.fetchall()
if not result:
return None
- res = result[0]
-
- for k in res:
- if isinstance(res[k], datetime):
- res[k] = res[k].replace(tzinfo=UTC)
-
- return GRToken.model_validate(res)
+ return GRToken.model_validate(result[0])
diff --git a/generalresearch/models/gr/authentication.py b/generalresearch/models/gr/authentication.py
index 5366db8..3c7c461 100644
--- a/generalresearch/models/gr/authentication.py
+++ b/generalresearch/models/gr/authentication.py
@@ -7,17 +7,15 @@ from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any
from pydantic import (
- AnyHttpUrl,
BaseModel,
ConfigDict,
EmailStr,
Field,
- NonNegativeInt,
PositiveInt,
field_validator,
)
-from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr
+from generalresearch.models.custom_types import AwareDatetimeISO, HttpsUrlStr, UUIDStr
if TYPE_CHECKING:
from generalresearch.managers.gr.business import BusinessManager
@@ -30,76 +28,10 @@ if TYPE_CHECKING:
from generalresearch.redis_helper import RedisConfig
-class Claims(BaseModel):
- iss: str | None = Field(
- default=None,
- description="Issuer: https://www.rfc-editor.org/rfc/rfc7519.html#section-4.1.1",
- )
-
- sub: str | None = Field(
- default=None,
- description="Subject: https://www.rfc-editor.org/rfc/rfc7519.html#section-4.1.2",
- )
-
- aud: str | None = Field(
- default=None,
- description="Audience: https://www.rfc-editor.org/rfc/rfc7519.html#section-4.1.3",
- )
-
- exp: NonNegativeInt | None = Field(
- default=None,
- description="Expiration time: https://www.rfc-editor.org/rfc/rfc7519.html#section-4.1.4",
- )
-
- iat: NonNegativeInt | None = Field(
- default=None,
- description="Issued at: https://www.rfc-editor.org/rfc/rfc7519.html#section-4.1.6",
- )
-
- auth_time: NonNegativeInt | None = Field(
- default=None,
- description="When authentication occured: https://openid.net/specs/openid-connect-core-1_0.html#IDToken",
- )
-
- acr: str | None = Field(
- default=None,
- description="Authentication Context Class Reference: https://openid.net/specs/openid-connect-core-1_0.html#IDToken",
- )
-
- amr: list[str] | None = Field(
- default=None,
- description="Authentication Methods References: https://openid.net/specs/openid-connect-core-1_0.html#IDToken",
- )
-
- c_hash: str | None = Field(
- default=None,
- description="Code hash value: http://openid.net/specs/openid-connect-core-1_0.html",
- )
-
- nonce: str | None = Field(
- default=None,
- description="Value used to associate a Client session with an ID Token: http://openid.net/specs/openid-connect-core-1_0.html",
- )
-
- at_hash: str | None = Field(
- default=None,
- description="Access Token hash value: http://openid.net/specs/openid-connect-core-1_0.html",
- )
-
- sid: str | None = Field(
- default=None,
- description="Session ID: https://openid.net/specs/openid-connect-frontchannel-1_0.html#ClaimsContents",
- )
-
- # --- Properties ---
-
- @property
- def subject(self):
- return self.sub
-
-
class GRUserRef(BaseModel):
email: EmailStr | None = Field(default=None)
+ name: str | None = Field(default=None)
+ avatar_url: HttpsUrlStr | None = Field(default=None)
date_joined: AwareDatetimeISO = Field(
description="When the GR User account signed up."
)
@@ -113,8 +45,10 @@ class GRUser(BaseModel):
)
id: PositiveInt | None = Field(default=None)
- sub: str | None = Field(max_length=200)
email: EmailStr | None = Field(default=None)
+ name: str | None = Field(default=None)
+ avatar_url: HttpsUrlStr | None = Field(default=None)
+
is_superuser: bool = Field(default=False)
date_joined: AwareDatetimeISO = Field(
@@ -126,28 +60,15 @@ class GRUser(BaseModel):
teams: list[Team] | None = Field(default=None)
products: list[Product] | None = Field(default=None)
token: GRToken | None = Field(default=None)
- claims: Claims | None = Field(default=None)
def to_user_ref(self) -> GRUserRef:
return GRUserRef(
email=self.email,
+ name=self.name,
+ avatar_url=self.avatar_url,
date_joined=self.date_joined,
)
- def prefetch_claims(
- self, token: str, key: dict[str, Any], audience: str, issuer: AnyHttpUrl
- ) -> None:
- from jose import jwt
-
- payload = jwt.decode(
- token=token,
- key=key,
- algorithms=["RS256"],
- audience=audience,
- issuer=issuer,
- )
- self.claims = Claims.model_validate(payload)
-
def prefetch_businesses(
self, pg_config: PostgresConfig, redis_config: RedisConfig
) -> None:
@@ -158,6 +79,7 @@ class GRUser(BaseModel):
if self.is_superuser:
self.businesses = bm.get_all()
else:
+ assert self.id is not None
self.businesses = bm.get_by_user_id(user_id=self.id)
def prefetch_teams(
@@ -206,6 +128,7 @@ class GRUser(BaseModel):
)
tm = GRTokenManager(pg_config=pg_config)
+ assert self.id is not None
self.token = tm.get_by_user_id(user_id=self.id)
def __eq__(self, other: GRUser) -> bool:
@@ -242,9 +165,9 @@ class GRUser(BaseModel):
team_manager: TeamManager | None,
product_manager: ProductManager,
) -> list[UUIDStr]:
- assert (
- business_manager is not None or team_manager is not None
- ), "At least one of business_manager or team_manager is required"
+ assert business_manager is not None or team_manager is not None, (
+ "At least one of business_manager or team_manager is required"
+ )
return product_manager.get_uuids(
business_uuids=(
diff --git a/test_utils/models/gr/conftest.py b/test_utils/models/gr/conftest.py
index d99bd89..4b9ff4c 100644
--- a/test_utils/models/gr/conftest.py
+++ b/test_utils/models/gr/conftest.py
@@ -5,6 +5,7 @@ from random import randint
from typing import TYPE_CHECKING
from uuid import uuid4
+import faker
import pytest
from pydantic import PositiveInt
from pydantic_extra_types.phone_numbers import PhoneNumber
@@ -28,7 +29,9 @@ if TYPE_CHECKING:
)
from generalresearch.models.gr.team import Membership, Team
from generalresearch.pg_helper import PostgresConfig
- from generalresearch.redis_helper import RedisConfig
+
+
+fake = faker.Faker()
# --- Static ---
@@ -216,15 +219,16 @@ def gr_user_factory(gr_user_manager: GRUserManager) -> Callable[..., GRUser]:
def _inner(
save: bool = True,
- sub: str | None = None,
+ name: str | None = None,
email: str | None = None,
is_superuser: bool = False,
) -> GRUser:
- sub = sub or f"{uuid4().hex}-{uuid4().hex}"
+ name = name or fake.name()
+ email = email or f"{uuid4().hex}@example.com"
if save:
return gr_user_manager.create(
- sub=sub,
+ name=name,
email=email,
is_superuser=is_superuser,
)
@@ -235,19 +239,6 @@ def gr_user_factory(gr_user_manager: GRUserManager) -> Callable[..., GRUser]:
@pytest.fixture
-def gr_user_cache(
- gr_user: GRUser,
- gr_db: PostgresConfig,
- thl_web_rr: PostgresConfig,
- gr_redis_config: RedisConfig,
-) -> GRUser:
- gr_user.set_cache(
- pg_config=gr_db, thl_web_rr=thl_web_rr, redis_config=gr_redis_config
- )
- return gr_user
-
-
-@pytest.fixture
def gr_user(gr_user_factory: Callable[..., GRUser]) -> GRUser:
return gr_user_factory(save=True)
diff --git a/tests/managers/gr/test_authentication.py b/tests/managers/gr/test_authentication.py
index 1310c79..324bfa6 100644
--- a/tests/managers/gr/test_authentication.py
+++ b/tests/managers/gr/test_authentication.py
@@ -1,11 +1,9 @@
-import logging
from collections.abc import Callable
from uuid import uuid4
import pytest
from generalresearch.managers.gr.authentication import GRTokenManager, GRUserManager
-from generalresearch.managers.gr.team import TeamManager
from generalresearch.models.gr.authentication import GRToken, GRUser
from generalresearch.pg_helper import PostgresConfig
from generalresearch.redis_helper import RedisConfig
@@ -14,12 +12,13 @@ SSO_ISSUER = ""
class TestGRUserManager:
-
def test_create(self, gr_user: GRUser, gr_user_manager: GRUserManager):
+ assert gr_user.id is not None
instance = gr_user_manager.get_by_id(gr_user.id)
assert isinstance(instance, GRUser)
assert gr_user.id == instance.id
+ assert gr_user.id is not None
instance2 = gr_user_manager.get_by_id(gr_user.id)
assert isinstance(instance2, GRUser)
assert gr_user.model_dump_json() == instance2.model_dump_json()
@@ -28,32 +27,25 @@ class TestGRUserManager:
with pytest.raises(expected_exception=ValueError) as cm:
gr_user_manager.get_by_id(gr_user_id=999_999_999)
assert "GRUser not found" in str(cm.value)
-
+ assert gr_user.id is not None
instance = gr_user_manager.get_by_id(gr_user_id=gr_user.id)
assert isinstance(instance, GRUser)
- assert instance.sub == gr_user.sub
-
- def test_get_by_sub(self, gr_user: GRUser, gr_user_manager: GRUserManager):
- with pytest.raises(expected_exception=ValueError) as cm:
- gr_user_manager.get_by_sub(sub=uuid4().hex)
- assert "GRUser not found" in str(cm.value)
-
- instance = gr_user_manager.get_by_sub(sub=gr_user.sub)
- assert isinstance(instance, GRUser)
- assert instance.id == gr_user.id
+ assert instance.email == gr_user.email
- def test_get_by_sub_or_create(
- self, gr_user: GRUser, gr_user_manager: GRUserManager
+ def test_get_by_email(
+ self,
+ gr_user_factory: Callable[..., GRUser],
+ gr_user_manager: GRUserManager,
):
- sub = f"{uuid4().hex}-{uuid4().hex}"
+ email = f"{uuid4().hex}@example.com"
+ gr_user = gr_user_factory(email=email)
- with pytest.raises(expected_exception=ValueError) as cm:
- gr_user_manager.get_by_sub(sub=sub)
- assert "GRUser not found" in str(cm.value)
+ with pytest.raises(expected_exception=ValueError, match="GRUser not found"):
+ gr_user_manager.get_by_email(email="missing@example.com")
- instance = gr_user_manager.get_by_sub_or_create(sub=sub)
+ instance = gr_user_manager.get_by_email(email=email)
assert isinstance(instance, GRUser)
- assert instance.sub == sub
+ assert instance.id == gr_user.id
def test_get_all(
self, gr_user_factory: Callable[..., GRUser], gr_user_manager: GRUserManager
@@ -70,21 +62,10 @@ class TestGRUserManager:
assert isinstance(res, list)
assert res == []
- def test_list_product_uuids(
- self,
- caplog,
- gr_user: GRUser,
- gr_user_manager: GRUserManager,
- thl_web_rr: PostgresConfig,
- ):
- with caplog.at_level(logging.WARNING):
- gr_user_manager.list_product_uuids(user=gr_user, thl_pg_config=thl_web_rr)
- assert "prefetch not run" in caplog.text
-
class TestGRTokenManager:
-
def test_create(self, gr_user: GRUser, gr_token_manager: GRTokenManager):
+ assert gr_user.id is not None
assert gr_token_manager.create(user_id=gr_user.id) is None
token = gr_token_manager.get_by_user_id(user_id=gr_user.id)
@@ -92,6 +73,7 @@ class TestGRTokenManager:
assert gr_user.id == token.user_id
def test_get_by_user_id(self, gr_user: GRUser, gr_token_manager: GRTokenManager):
+ assert gr_user.id is not None
assert gr_token_manager.create(user_id=gr_user.id) is None
token = gr_token_manager.get_by_user_id(user_id=gr_user.id)
@@ -105,7 +87,7 @@ class TestGRTokenManager:
gr_db: PostgresConfig,
gr_redis_config: RedisConfig,
):
-
+ assert gr_user.id is not None
gr_token_manager.create(user_id=gr_user.id)
token: GRToken | None = gr_token_manager.get_by_user_id(user_id=gr_user.id)
@@ -120,6 +102,7 @@ class TestGRTokenManager:
gr_user: GRUser,
gr_token_manager: GRTokenManager,
):
+ assert gr_user.id is not None
gr_token_manager.create(user_id=gr_user.id)
token = gr_token_manager.get_by_user_id(user_id=gr_user.id)
assert isinstance(token, GRToken)
@@ -131,25 +114,3 @@ class TestGRTokenManager:
with pytest.raises(expected_exception=Exception) as cm:
gr_token_manager.get_by_key(api_key=uuid4().hex)
assert "No GRUser with token of " in str(cm.value)
-
- @pytest.mark.skip(reason="no idea how to actually test this...")
- def test_get_by_sso_key(
- self,
- gr_team_manager: TeamManager,
- gr_redis_config: RedisConfig,
- ):
-
- api_key = "..."
- jwks = {
- # ...
- }
-
- instance = gr_team_manager.get_by_key(
- api_key=api_key,
- jwks=jwks,
- audience="...",
- issuer=SSO_ISSUER,
- gr_redis_config=gr_redis_config,
- )
-
- assert isinstance(instance, GRToken)
diff --git a/tests/managers/gr/test_team.py b/tests/managers/gr/test_team.py
index 689c4be..6c38a49 100644
--- a/tests/managers/gr/test_team.py
+++ b/tests/managers/gr/test_team.py
@@ -11,13 +11,13 @@ if TYPE_CHECKING:
from generalresearch.managers.gr.authentication import GRUserManager
from generalresearch.managers.gr.team import MembershipManager, TeamManager
from generalresearch.models.gr.authentication import GRUser
+ from generalresearch.models.gr.business import Business
from generalresearch.models.thl.product import Product
from generalresearch.pg_helper import PostgresConfig
from generalresearch.redis_helper import RedisConfig
class TestMembershipManager:
-
def test_init(
self, gr_membership_manager: MembershipManager, gr_db: PostgresConfig
):
@@ -25,7 +25,6 @@ class TestMembershipManager:
class TestTeamManager:
-
def test_init(self, gr_team_manager: TeamManager, gr_db: PostgresConfig):
assert gr_team_manager.pg_config == gr_db
@@ -81,13 +80,11 @@ class TestTeamManager:
gr_team_manager: TeamManager,
gr_user_manager: GRUserManager,
gr_user_factory: Callable[..., GRUser],
+ gr_membership_factory: Callable[..., Membership],
):
user: GRUser = gr_user_factory()
-
- instance = gr_team_manager.add_user(
- gr_user_manager=gr_user_manager, team=gr_team, gr_user=user
- )
+ instance = gr_membership_factory(gr_team=gr_team, gr_user=user)
assert isinstance(instance, Membership)
# assert team.gr_users is None
@@ -122,11 +119,10 @@ class TestTeamManager:
gr_user_factory: Callable[..., GRUser],
gr_team_manager: TeamManager,
gr_user_manager: GRUserManager,
+ gr_membership_factory: Callable[..., Membership],
):
user: GRUser = gr_user_factory()
- gr_team_manager.add_user(
- gr_user_manager=gr_user_manager, team=gr_team, gr_user=user
- )
+ gr_membership_factory(gr_team=gr_team, gr_user=user)
res = gr_team_manager.get_by_user(gr_user=user)
assert isinstance(res, list)
@@ -135,6 +131,35 @@ class TestTeamManager:
assert isinstance(instance, Team)
assert instance.uuid == gr_team.uuid
+ def test_filter_includes_relationship_counts(
+ self,
+ gr_team_factory: Callable[..., Team],
+ gr_team_manager: TeamManager,
+ gr_user_factory: Callable[..., GRUser],
+ gr_membership_factory: Callable[..., Membership],
+ gr_business_factory: Callable[..., Business],
+ ):
+ team = gr_team_factory()
+ other_team = gr_team_factory()
+ user = gr_user_factory()
+ other_user = gr_user_factory()
+ business = gr_business_factory()
+
+ gr_membership_factory(gr_team=team, gr_user=user)
+ gr_membership_factory(gr_team=team, gr_user=other_user)
+ gr_team_manager.add_business(team_id=team.id, business_id=business.id)
+
+ assert gr_team_manager.filter(team_id=other_team.id) == [other_team]
+
+ teams = gr_team_manager.filter(
+ team_uuid=team.uuid,
+ gr_user_id=user.id,
+ )
+ assert len(teams) == 1
+ assert teams[0].id == team.id
+ assert teams[0].business_count == 1
+ assert teams[0].membership_count == 2
+
def test_get_by_user_duplicates(
self,
gr_user: GRUser,