diff options
| -rw-r--r-- | generalresearch/gr_django/migrations/0004_user_avatar_url_user_name_alter_user_sub.py | 28 | ||||
| -rw-r--r-- | generalresearch/gr_django/models.py | 8 | ||||
| -rw-r--r-- | generalresearch/managers/gr/authentication.py | 163 | ||||
| -rw-r--r-- | generalresearch/models/gr/authentication.py | 103 | ||||
| -rw-r--r-- | pyproject.toml | 2 | ||||
| -rw-r--r-- | test_utils/models/gr/conftest.py | 25 | ||||
| -rw-r--r-- | tests/managers/gr/test_authentication.py | 75 | ||||
| -rw-r--r-- | tests/managers/gr/test_team.py | 43 |
8 files changed, 157 insertions, 290 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..54ebaff 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 | None = None, + 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"] @@ -58,11 +59,31 @@ class GRUserManager(PostgresManagerWithRedis): instance.id = gr_user_id return instance + def update(self, user: GRUser) -> None: + """Persist the user's editable profile fields.""" + assert user.id is not None + data = user.model_dump( + mode="json", include={"id", "name", "avatar_url"} + ) + + with self.pg_config.make_connection() as conn: + with conn.cursor() as c: + c.execute( + query=sql.SQL(""" + UPDATE gr_user + SET name = %(name)s, avatar_url = %(avatar_url)s + WHERE id = %(id)s + """), + params=data, + ) + if c.rowcount == 0: + raise ValueError("GRUser not found") + conn.commit() + def filter( 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 +93,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 +126,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 +138,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 +168,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 +187,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 +230,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/pyproject.toml b/pyproject.toml index 3e86956..c4c4e29 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "generalresearch" -version = "3.9.5" +version = "3.9.6" description = "Python Utilities for General Research" readme = "README.md" requires-python = ">=3.14" 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, |