diff options
| -rw-r--r-- | generalresearch/grliq/managers/forensic_data.py | 330 | ||||
| -rw-r--r-- | generalresearch/incite/base.py | 170 | ||||
| -rw-r--r-- | generalresearch/managers/gr/authentication.py | 79 | ||||
| -rw-r--r-- | generalresearch/managers/gr/business.py | 175 | ||||
| -rw-r--r-- | generalresearch/managers/gr/team.py | 93 | ||||
| -rw-r--r-- | generalresearch/managers/leaderboard/__init__.py | 4 | ||||
| -rw-r--r-- | generalresearch/managers/thl/ipinfo.py | 106 | ||||
| -rw-r--r-- | pyproject.toml | 1 | ||||
| -rw-r--r-- | test_utils/managers/gr/__init__.py | 0 | ||||
| -rw-r--r-- | test_utils/managers/gr/conftest.py | 151 | ||||
| -rw-r--r-- | test_utils/managers/grliq/__init__.py | 0 | ||||
| -rw-r--r-- | test_utils/managers/grliq/conftest.py | 61 | ||||
| -rw-r--r-- | test_utils/managers/thl/__init__.py | 0 | ||||
| -rw-r--r-- | test_utils/managers/thl/conftest.py | 112 | ||||
| -rw-r--r-- | test_utils/models/conftest.py | 201 | ||||
| -rw-r--r-- | tests/incite/test_collection_base.py | 6 | ||||
| -rw-r--r-- | tests/models/custom_types/test_aware_datetime.py | 5 | ||||
| -rw-r--r-- | tests/models/thl/test_product.py | 53 |
18 files changed, 811 insertions, 736 deletions
diff --git a/generalresearch/grliq/managers/forensic_data.py b/generalresearch/grliq/managers/forensic_data.py index d7e362d..739c520 100644 --- a/generalresearch/grliq/managers/forensic_data.py +++ b/generalresearch/grliq/managers/forensic_data.py @@ -1,11 +1,11 @@ -from datetime import datetime, timezone -from typing import Any, Collection, Dict, List, Optional, Tuple -from uuid import uuid4 +from __future__ import annotations + +from datetime import datetime +from typing import Any, Collection from psycopg import sql from pydantic import NonNegativeInt, PositiveInt -from generalresearch.grliq.managers import DUMMY_GRLIQ_DATA from generalresearch.grliq.models.events import PointerMove, TimingData from generalresearch.grliq.models.forensic_data import GrlIqData from generalresearch.grliq.models.forensic_result import ( @@ -23,58 +23,13 @@ class GrlIqDataManager: def __init__(self, postgres_config: PostgresConfig): self.postgres_config = postgres_config - def create_dummy( - self, - is_attempt_allowed: bool = True, - product_id: Optional[str] = None, - product_user_id: Optional[str] = None, - uuid: Optional[str] = None, - mid: Optional[str] = None, - created_at: Optional[datetime] = None, - ) -> GrlIqData: - """ - Creates a dummy record in the db with a GrlIqData (data), GrlIqCheckerResults (result_data), - and GrlIqForensicCategoryResult (category_results) - :param is_attempt_allowed: Whether the attempt is allowed. - :param product_id: product_id of user - :param product_user_id: product_user_id of user - :param uuid: uuid for the grliq data record - :param mid: the thl_session:uuid / mid for the attempt. - :return: - """ - import copy - - res: GrlIqData = copy.deepcopy(DUMMY_GRLIQ_DATA[int(is_attempt_allowed)]) - - product_id = product_id or uuid4().hex - product_user_id = product_user_id or uuid4().hex - uuid = uuid or uuid4().hex - mid = mid or uuid4().hex - created_at = created_at or datetime.now(tz=timezone.utc) - - res["data"].product_id = product_id - res["data"].product_user_id = product_user_id - res["data"].uuid = uuid - res["data"].mid = mid - res["data"].created_at = created_at - res["result_data"].uuid = uuid - res["category_result"].uuid = uuid - - return self.create( - iq_data=res["data"], - result_data=res["result_data"], - category_result=res["category_result"], - fraud_score=res["category_result"].fraud_score, - is_attempt_allowed=res["category_result"].is_attempt_allowed(), - ) - def create( self, iq_data: GrlIqData, - result_data: Optional[GrlIqCheckerResults] = None, - category_result: Optional[GrlIqForensicCategoryResult] = None, - fraud_score: Optional[int] = None, - is_attempt_allowed: Optional[bool] = None, + result_data: GrlIqCheckerResults | None = None, + category_result: GrlIqForensicCategoryResult | None = None, + fraud_score: int | None = None, + is_attempt_allowed: bool | None = None, ) -> GrlIqData: data = iq_data.model_dump_sql(exclude={"events", "mouse_events", "timing_data"}) @@ -95,8 +50,7 @@ class GrlIqDataManager: data["fraud_score"] = fraud_score data["is_attempt_allowed"] = is_attempt_allowed - query = sql.SQL( - """ + query = sql.SQL(""" INSERT INTO grliq_forensicdata (uuid, session_uuid, created_at, product_id, product_user_id, country_iso, client_ip, ua_browser_family, ua_browser_version, @@ -112,14 +66,12 @@ class GrlIqDataManager: %(fingerprint)s, %(fraud_score)s, %(is_attempt_allowed)s, %(result_data)s, %(category_result)s) RETURNING id - """ - ) + """) - with self.postgres_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, data) - pk = c.fetchone()["id"] # type: ignore - conn.commit() + with self.postgres_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, data) + pk = c.fetchone()["id"] # type: ignore + conn.commit() iq_data.id = pk @@ -130,9 +82,9 @@ class GrlIqDataManager: uuid: UUIDStr, result_data: GrlIqCheckerResults, category_result: GrlIqForensicCategoryResult, - fingerprint: Optional[str] = None, - fraud_score: Optional[int] = None, - is_attempt_allowed: Optional[bool] = None, + fingerprint: str | None = None, + fraud_score: int | None = None, + is_attempt_allowed: bool | None = None, ) -> None: data = {"uuid": uuid} data["result_data"] = result_data.model_dump_json(exclude_none=True) @@ -141,8 +93,7 @@ class GrlIqDataManager: data["fraud_score"] = fraud_score data["is_attempt_allowed"] = is_attempt_allowed - query = sql.SQL( - """ + query = sql.SQL(""" UPDATE grliq_forensicdata SET result_data = %(result_data)s, category_result = %(category_result)s, @@ -150,8 +101,7 @@ class GrlIqDataManager: fraud_score = %(fraud_score)s, is_attempt_allowed = %(is_attempt_allowed)s WHERE uuid = %(uuid)s - """ - ) + """) with self.postgres_config.make_connection() as conn: with conn.cursor() as c: c.execute(query, data) @@ -161,21 +111,17 @@ class GrlIqDataManager: ) conn.commit() - return None - def update_fingerprint(self, iq_data: GrlIqData) -> None: # We should only run this if we modified the fingerprint algorithm if "fingerprint" in iq_data.__dict__: # make sure it's not cached del iq_data.__dict__["fingerprint"] data = {"uuid": iq_data.uuid, "fingerprint": iq_data.fingerprint} - query = sql.SQL( - """ + query = sql.SQL(""" UPDATE grliq_forensicdata SET fingerprint = %(fingerprint)s WHERE uuid = %(uuid)s - """ - ) + """) with self.postgres_config.make_connection() as conn: with conn.cursor() as c: c.execute(query, data) @@ -189,13 +135,11 @@ class GrlIqDataManager: # We should only run this if we structured new fields and want to # back-populate them in the db data = {"id": iq_data.id, "data": iq_data.model_dump_sql()["data"]} - query = sql.SQL( - """ + query = sql.SQL(""" UPDATE grliq_forensicdata SET data = %(data)s WHERE id = %(id)s - """ - ) + """) with self.postgres_config.make_connection() as conn: with conn.cursor() as c: c.execute(query, data) @@ -207,7 +151,7 @@ class GrlIqDataManager: def get_data_if_exists( self, forensic_uuid: UUIDStr, load_events: bool = False - ) -> Optional[GrlIqData]: + ) -> GrlIqData | None: try: return self.get_data(forensic_uuid=forensic_uuid, load_events=load_events) except AssertionError: @@ -215,8 +159,8 @@ class GrlIqDataManager: def get_data( self, - forensic_id: Optional[PositiveInt] = None, - forensic_uuid: Optional[UUIDStr] = None, + forensic_id: PositiveInt | None = None, + forensic_uuid: UUIDStr | None = None, load_events: bool = False, ) -> GrlIqData: from generalresearch.grliq.managers.forensic_events import ( @@ -230,8 +174,7 @@ class GrlIqDataManager: # forensic items' session, 2) event_start is closest to the # created_at for this forensic item, and within 1 minute. - query = sql.SQL( - """ + query = sql.SQL(""" SELECT d.id, d.data, e.events, e.mouse_events, t.timing_data FROM grliq_forensicdata d -- Closest event_start within 1 minute @@ -251,16 +194,13 @@ class GrlIqDataManager: ORDER BY e2.id DESC LIMIT 1 ) t ON true - """ - ) + """) else: - query = sql.SQL( - """ + query = sql.SQL(""" SELECT d.id, d.data FROM grliq_forensicdata d - """ - ) + """) if forensic_id is not None: column_name = "id" @@ -275,10 +215,9 @@ class GrlIqDataManager: limit_clause = sql.SQL(" LIMIT 1") q1 = sql.Composed([query, where_clause, limit_clause]) - with self.postgres_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query=q1, params=(param_value,)) - x = c.fetchone() + with self.postgres_config.make_connection() as conn, conn.cursor() as c: + c.execute(query=q1, params=(param_value,)) + x = c.fetchone() assert x is not None, f"GrlIqDataManager.get_data({forensic_uuid=}) not found" @@ -316,10 +255,10 @@ class GrlIqDataManager: def filter_timing_data( self, - created_between: Tuple[datetime, datetime], - limit: Optional[int] = None, - offset: Optional[int] = None, - ) -> List[Dict[str, Any]]: + created_between: tuple[datetime, datetime], + limit: int | None = None, + offset: int | None = None, + ) -> list[dict[str, Any]]: # TODO! created_between used to be marked as Optional, but it would # break the query. Evaluate it's use to determine best behavior. @@ -349,10 +288,9 @@ class GrlIqDataManager: {limit_str} {offset_str}; """ - with self.postgres_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, params) - res: List[Dict[str, Any]] = c.fetchall() # type: ignore + with self.postgres_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, params) + res: list[dict[str, Any]] = c.fetchall() # type: ignore for x in res: x["timing_data"] = TimingData.model_validate(x["timing_data"]) @@ -369,47 +307,44 @@ class GrlIqDataManager: # This is used for filtering for other forensic posts with a certain # fingerprint, in this product_id, but NOT for this user. - query = sql.SQL( - """ + query = sql.SQL(""" SELECT COUNT(DISTINCT product_user_id) as user_count FROM grliq_forensicdata d WHERE product_id = %(product_id)s AND fingerprint = %(fingerprint)s AND product_user_id != %(product_user_id)s AND created_at > NOW() - INTERVAL '30 DAYS' - """ - ) + """) params = { "product_id": product_id, "fingerprint": fingerprint, "product_user_id": product_user_id_not, } # print(query) - with self.postgres_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query, params) - user_count = c.fetchone()["user_count"] # type: ignore + with self.postgres_config.make_connection() as conn, conn.cursor() as c: + c.execute(query, params) + user_count = c.fetchone()["user_count"] # type: ignore return int(user_count) def filter_data( self, - session_uuid: Optional[str] = None, - fingerprint: Optional[str] = None, - fingerprints: Optional[Collection[str]] = None, - product_id: Optional[str] = None, - product_ids: Optional[Collection[str]] = None, - uuids: Optional[Collection[str]] = None, - created_after: Optional[datetime] = None, - created_before: Optional[datetime] = None, - created_between: Optional[Tuple[datetime, datetime]] = None, - user: Optional[User] = None, - users: Optional[Collection[User]] = None, - phase: Optional[Phase] = None, + session_uuid: str | None = None, + fingerprint: str | None = None, + fingerprints: Collection[str] | None = None, + product_id: str | None = None, + product_ids: Collection[str] | None = None, + uuids: Collection[str] | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, + created_between: tuple[datetime, datetime] | None = None, + user: User | None = None, + users: Collection[User] | None = None, + phase: Phase | None = None, order_by: str = "created_at DESC", - limit: Optional[int] = None, - offset: Optional[int] = None, - ) -> List[GrlIqData]: + limit: int | None = None, + offset: int | None = None, + ) -> list[GrlIqData]: res = self.filter( select_str="d.id, d.data", @@ -433,18 +368,18 @@ class GrlIqDataManager: def filter_results( self, - session_uuid: Optional[str] = None, - uuid: Optional[str] = None, - product_ids: Optional[Collection[str]] = None, - product_id: Optional[str] = None, - created_after: Optional[datetime] = None, - created_before: Optional[datetime] = None, - created_between: Optional[Tuple[datetime, datetime]] = None, - user: Optional[User] = None, - limit: Optional[int] = None, - offset: Optional[int] = None, + session_uuid: str | None = None, + uuid: str | None = None, + product_ids: Collection[str] | None = None, + product_id: str | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, + created_between: tuple[datetime, datetime] | None = None, + user: User | None = None, + limit: int | None = None, + offset: int | None = None, order_by: str = "created_at DESC", - ) -> List[GrlIqCheckerResults]: + ) -> list[GrlIqCheckerResults]: select_str = ( "id, session_uuid, product_id, product_user_id, created_at, result_data" ) @@ -472,18 +407,18 @@ class GrlIqDataManager: def filter_category_results( self, - session_uuid: Optional[str] = None, - uuid: Optional[str] = None, - product_id: Optional[str] = None, - product_ids: Optional[Collection[str]] = None, - created_after: Optional[datetime] = None, - created_before: Optional[datetime] = None, - created_between: Optional[Tuple[datetime, datetime]] = None, - user: Optional[User] = None, + session_uuid: str | None = None, + uuid: str | None = None, + product_id: str | None = None, + product_ids: Collection[str] | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, + created_between: tuple[datetime, datetime] | None = None, + user: User | None = None, order_by: str = "created_at DESC", - limit: Optional[int] = None, - offset: Optional[int] = None, - ) -> List[GrlIqForensicCategoryResult]: + limit: int | None = None, + offset: int | None = None, + ) -> list[GrlIqForensicCategoryResult]: select_str = ( "id, session_uuid, product_id, product_user_id, created_at, category_result" ) @@ -506,22 +441,22 @@ class GrlIqDataManager: @staticmethod def make_filter_str( - session_uuid: Optional[str] = None, - fingerprint: Optional[str] = None, - fingerprints: Optional[Collection[str]] = None, - uuids: Optional[Collection[str]] = None, - product_id: Optional[str] = None, - product_ids: Optional[Collection[str]] = None, - created_after: Optional[datetime] = None, - created_before: Optional[datetime] = None, - created_between: Optional[Tuple[datetime, datetime]] = None, - user: Optional[User] = None, - users: Optional[Collection[User]] = None, - phase: Optional[Phase] = None, - ) -> Tuple[str, Dict[str, Any]]: + session_uuid: str | None = None, + fingerprint: str | None = None, + fingerprints: Collection[str] | None = None, + uuids: Collection[str] | None = None, + product_id: str | None = None, + product_ids: Collection[str] | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, + created_between: tuple[datetime, datetime] | None = None, + user: User | None = None, + users: Collection[User] | None = None, + phase: Phase | None = None, + ) -> tuple[str, dict[str, Any]]: filters = [] - params: Dict[str, Any] = {} + params: dict[str, Any] = {} if session_uuid: params["session_uuid"] = session_uuid @@ -614,19 +549,20 @@ class GrlIqDataManager: def filter_count( self, - session_uuid: Optional[str] = None, - fingerprint: Optional[str] = None, - fingerprints: Optional[Collection[str]] = None, - uuids: Optional[Collection[str]] = None, - product_id: Optional[str] = None, - product_ids: Optional[Collection[str]] = None, - created_after: Optional[datetime] = None, - created_before: Optional[datetime] = None, - created_between: Optional[Tuple[datetime, datetime]] = None, - user: Optional[User] = None, - users: Optional[Collection[User]] = None, - phase: Optional[Phase] = None, + session_uuid: str | None = None, + fingerprint: str | None = None, + fingerprints: Collection[str] | None = None, + uuids: Collection[str] | None = None, + product_id: str | None = None, + product_ids: Collection[str] | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, + created_between: tuple[datetime, datetime] | None = None, + user: User | None = None, + users: Collection[User] | None = None, + phase: Phase | None = None, ) -> NonNegativeInt: + filter_str, params = self.make_filter_str( session_uuid=session_uuid, fingerprint=fingerprint, @@ -682,36 +618,35 @@ class GrlIqDataManager: FROM grliq_forensicdata d {filter_str} """ - with self.postgres_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query=query, params=params) - res = c.fetchone() + with self.postgres_config.make_connection() as conn, conn.cursor() as c: + c.execute(query=query, params=params) + res = c.fetchone() return int(res["c"]) def filter( self, select_str: str, - session_uuid: Optional[str] = None, - fingerprint: Optional[str] = None, - fingerprints: Optional[Collection[str]] = None, - uuids: Optional[Collection[str]] = None, - product_id: Optional[str] = None, - product_ids: Optional[Collection[str]] = None, - created_after: Optional[datetime] = None, - created_before: Optional[datetime] = None, - created_between: Optional[Tuple[datetime, datetime]] = None, - user: Optional[User] = None, - users: Optional[Collection[User]] = None, - phase: Optional[Phase] = None, + session_uuid: str | None = None, + fingerprint: str | None = None, + fingerprints: Collection[str] | None = None, + uuids: Collection[str] | None = None, + product_id: str | None = None, + product_ids: Collection[str] | None = None, + created_after: datetime | None = None, + created_before: datetime | None = None, + created_between: tuple[datetime, datetime] | None = None, + user: User | None = None, + users: Collection[User] | None = None, + phase: Phase | None = None, order_by: str = "created_at DESC", - limit: Optional[int] = None, - offset: Optional[int] = None, - ) -> List[Dict[str, Any]]: + limit: int | None = None, + offset: int | None = None, + ) -> list[dict[str, Any]]: """ Accepts lots of optional filters. """ if not limit: - limit = 5000 + limit = 5_000 if not offset: offset = 0 @@ -746,10 +681,9 @@ class GrlIqDataManager: OFFSET {offset} """ # print(query) - with self.postgres_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query=query, params=params) - res: List[Dict[str, Any]] = c.fetchall() # type: ignore + with self.postgres_config.make_connection() as conn, conn.cursor() as c: + c.execute(query=query, params=params) + res: list[dict[str, Any]] = c.fetchall() # type: ignore for x in res: @@ -777,7 +711,7 @@ class GrlIqDataManager: return res @staticmethod - def temporary_add_missing_fields(d: Dict[str, Any]) -> None: + def temporary_add_missing_fields(d: dict[str, Any]) -> None: # The following fields were added recently, and so we must give them # a value or old db rows won't be parseable. Once logs are backfilled # then this can be removed diff --git a/generalresearch/incite/base.py b/generalresearch/incite/base.py index aa64bf0..a8088ac 100644 --- a/generalresearch/incite/base.py +++ b/generalresearch/incite/base.py @@ -18,11 +18,7 @@ from typing import ( TYPE_CHECKING, Any, Callable, - List, - Optional, Sequence, - Tuple, - Union, ) from uuid import uuid4 @@ -30,7 +26,7 @@ import dask import dask.dataframe as dd import pandas as pd import pyarrow.parquet as pq -from distributed import Client +from distributed import Client as DaskClient from pandera.pandas import DataFrameSchema from pydantic import ( BaseModel, @@ -38,7 +34,9 @@ from pydantic import ( DirectoryPath, Field, FilePath, + PositiveInt, PrivateAttr, + TypeAdapter, ValidationInfo, field_validator, model_validator, @@ -61,7 +59,7 @@ if TYPE_CHECKING: ) from generalresearch.incite.mergers import MergeCollection, MergeType - Collection = Union[DFCollection, MergeCollection] + Collection = DFCollection | MergeCollection logging.basicConfig() LOG = logging.getLogger() @@ -71,6 +69,9 @@ Item = Any Items = Sequence[Item] DT_STR = "%Y-%m-%d %H:%M:%S" +_dir_adapter = TypeAdapter(DirectoryPath) +_filepath_adapter = TypeAdapter(FilePath) + class NFSMount(BaseModel): address: str = Field(default="127.0.0.1") @@ -89,8 +90,8 @@ class GRLDatasets(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) - data_src: Optional[Path] = Field(default=None) - incite: Optional[NFSMount] = Field(default=None) + data_src: Path | None = Field(default=None) + incite: NFSMount | None = Field(default=None) @model_validator(mode="after") def check_data_src_and_et_path(self) -> Self: @@ -99,6 +100,8 @@ class GRLDatasets(BaseModel): ) from generalresearch.incite.mergers import MergeType + assert self.data_src, "data src must be defined" + # Create the base folders and confirm we have read access self.data_src.mkdir(parents=True, exist_ok=True) assert access( @@ -121,7 +124,7 @@ class GRLDatasets(BaseModel): assert access(path=p, mode=R_OK), f"Cannot read {p}" return self - def archive_path(self, enum_type: Union[MergeType, DFCollectionType]) -> Path: + def archive_path(self, enum_type: MergeType | DFCollectionType) -> Path: """ TODO: Extend this so that it takes any type of Enum and that inputs in the correct parent dir for the respective Enum @@ -135,7 +138,7 @@ class GRLDatasets(BaseModel): pjoin(self.data_src, self.incite.point, folder, str(enum_type.value)) ) - def has_data(self, enum_type: Union[MergeType, DFCollectionType]) -> bool: + def has_data(self, enum_type: MergeType | DFCollectionType) -> bool: path_dir = self.archive_path(enum_type=enum_type) if isdir(path_dir): return bool(listdir(path_dir)) @@ -152,7 +155,7 @@ class CollectionBase(BaseModel): extra="forbid", ) - archive_path: DirectoryPath = Field(default="/tmp/") + archive_path: DirectoryPath = Field(default=_dir_adapter.validate_python("/tmp/")) df: SkipJsonSchema[pd.DataFrame] = Field( default_factory=lambda: pd.DataFrame(), exclude=True ) @@ -169,12 +172,12 @@ class CollectionBase(BaseModel): frozen=True, ) - finished: Optional[AwareDatetimeISO] = Field( + finished: AwareDatetimeISO | None = Field( default=None, description="Finished is only set if we don't want a rolling window", ) - _client: Optional[Client] = PrivateAttr(default=None) + _client: DaskClient | None = PrivateAttr(default=None) # --- Validators --- @model_validator(mode="before") @@ -188,15 +191,14 @@ class CollectionBase(BaseModel): assert isinstance(ap, Path), "check_model_before.isinstance(ap, Path)" if not ap.is_dir(): - raise ValueError(f"Path does not point to a directory") + raise ValueError("Path does not point to a directory") if not access(path=ap, mode=R_OK): - raise ValueError(f"Cannot read archive_path") + raise ValueError("Cannot read archive_path") - df: Optional[pd.DataFrame] = data.get("df", None) - if df is not None: - if not df.empty or len(df.columns) != 0: - raise ValueError("Do not provide a pd.DataFrame") + df: pd.DataFrame | None = data.get("df", None) + if df is not None and (not df.empty or len(df.columns) != 0): + raise ValueError("Do not provide a pd.DataFrame") return data @@ -215,21 +217,21 @@ class CollectionBase(BaseModel): @field_validator("start") def check_start( - cls, start: Optional[datetime], info: ValidationInfo - ) -> Optional[datetime]: + cls, start: datetime | None, info: ValidationInfo + ) -> datetime | None: if start and start.microsecond != 0: raise ValueError("Collection.start must not have microseconds") return start @field_validator("offset") - def check_offset(cls, v: Optional[str], info: ValidationInfo): + def check_offset(cls, v: str | None, info: ValidationInfo): # pd.offsets.__all__ if v is None: # In MergeCollections, offset can be None return v try: pd.Timedelta(v) - except (Exception,) as e: + except Exception as e: capture_exception(error=e) raise ValueError( "Invalid offset alias provided. Please review: " @@ -251,6 +253,7 @@ class CollectionBase(BaseModel): assert end, "an end value must be provided" _start = self.interval_start + assert _start, "a start value must be provided" if end.tzinfo is None: # A Naive end was passed in. We probably did this on purpose. @@ -283,13 +286,13 @@ class CollectionBase(BaseModel): ) @property - def interval_start(self) -> Optional[datetime]: + def interval_start(self) -> datetime | None: # In DFCollections, start must be set, so the interval_start = start. In merged # this may be overridden with different behavior. return self.start @property - def interval_range(self) -> List[Tuple]: + def interval_range(self) -> list[tuple[datetime, datetime]]: """closed='left', so 0 <= x < 5""" end = self.finished or datetime.now(tz=timezone.utc).replace(microsecond=0) iv_r = self._interval_range(end) @@ -302,7 +305,7 @@ class CollectionBase(BaseModel): return pd.DataFrame.from_records(records, index=self._interval_range(end)) @property - def items(self) -> pd.DataFrame: + def items(self) -> Items | None: raise NotImplementedError("Must override") @property @@ -315,19 +318,20 @@ class CollectionBase(BaseModel): def fetch_all_paths( self, - items: Optional[Items] = None, - force_rr_latest=False, - include_partial=False, - ) -> List[FilePath]: + items: Items | None = None, + force_rr_latest: bool = False, + include_partial: bool = False, + ) -> list[FilePath]: LOG.info( f"CollectionBase.fetch_all(items={len(items or [])}, " f"{force_rr_latest=}, {include_partial=})" ) items = items or self.items + assert items # (1) All the originally available archives - sources: List[FilePath] = [ + sources: list[FilePath] = [ i.path for i in items if i.has_archive(include_empty=False) ] @@ -357,14 +361,14 @@ class CollectionBase(BaseModel): def ddf( self, - items: Optional[Items] = None, - force_rr_latest=False, + items: Items | None = None, + force_rr_latest: bool = False, columns=None, filters=None, categories=None, include_partial=False, - graph: Optional[Callable] = None, - ) -> Optional[dd.DataFrame]: + graph: Callable | None = None, + ) -> dd.DataFrame | None: """ Args: @@ -396,7 +400,7 @@ class CollectionBase(BaseModel): """ if isinstance(items, list) and len(items): - sources: List[FilePath] = [ + sources: list[FilePath] = [ i.path for i in items if i.has_archive(include_empty=False) ] @@ -410,7 +414,7 @@ class CollectionBase(BaseModel): ) else: - sources: List[FilePath] = self.fetch_all_paths( + sources: list[FilePath] = self.fetch_all_paths( items=None, force_rr_latest=force_rr_latest, include_partial=include_partial, @@ -444,8 +448,8 @@ class CollectionBase(BaseModel): # --- Methods: Cleanup --- def schedule_cleanup( - self, client=None, sync=True, client_resources=None - ) -> Union[pd.DataFrame, Future]: + self, client: DaskClient | None = None, sync: bool = True, client_resources=None + ) -> pd.DataFrame | Future: LOG.info(f"cleanup(archive_path={self.archive_path})") fs = [] @@ -453,6 +457,8 @@ class CollectionBase(BaseModel): fs.append(dask.delayed(item.cleanup_partials)()) fs.append(dask.delayed(item.clear_corrupt_archive)()) fs.append(dask.delayed(self.clear_tmp_archives)()) + + assert isinstance(client, DaskClient) res = client.compute( collections=fs, sync=sync, @@ -468,15 +474,13 @@ class CollectionBase(BaseModel): self.clear_corrupt_archives() # self.check_empty() # what did this do?? - return None - def cleanup_partials(self) -> None: """If an item is "closed", remove any partial files that may be around...""" + assert self.items + for item in self.items: item.cleanup_partials() - return None - def clear_tmp_archives(self) -> None: regex = re.compile(r"\.parquet\.[0-9a-f]{32}", re.I) @@ -487,19 +491,18 @@ class CollectionBase(BaseModel): Path(os.path.join(self.archive_path, fn)) ) - return None - def clear_corrupt_archives(self) -> None: + assert self.items + for item in self.items: item.clear_corrupt_archive() - return None - def rebuild_symlinks(self) -> None: """ When copying "things" between filesystems, and using Sambda mmfsylinks, we can't ensure links are properly shared. """ + assert self.items for item in reversed(self.items): item: CollectionItemBase @@ -513,7 +516,6 @@ class CollectionBase(BaseModel): # Don't "continue" onto the next CollectionItem. Later on, # we may need to create a symlink for the most recent partial - pass # --- Empty Path --- if os.path.exists(empty_path): @@ -574,7 +576,7 @@ class CollectionBase(BaseModel): # `ln` command is run. -- Max 2024-07-26 try: os.remove(item.path.as_posix()) - except FileNotFoundError as e: + except FileNotFoundError: pass if platform == "darwin": @@ -582,16 +584,20 @@ class CollectionBase(BaseModel): else: subprocess.call(["ln", "-sfnT", highest_version, item.path.as_posix()]) - return None - # -- Methods: Source timing def get_item(self, interval: pd.Interval) -> Item: + assert self.items + return next(x for x in self.items if x.interval == interval) def get_item_start(self, start: pd.Timestamp) -> Items: + assert self.items + return next(x for x in self.items if x.interval.left == start) def get_items(self, since: datetime) -> Items: + assert self.items + res = [] first_match = True @@ -610,7 +616,7 @@ class CollectionBase(BaseModel): res.append(item) first_match = False - res: List[Item] = [i for i in res if not i.is_empty()] + res: list[Item] = [i for i in res if not i.is_empty()] if len([1 for i in res if i.should_archive() and not i.has_archive()]): warnings.warn( message="DFCollectionItem has missing archives", @@ -620,7 +626,7 @@ class CollectionBase(BaseModel): return res def get_items_from_year(self, year: int) -> Items: - ts = datetime(year=year, month=1, day=1) + ts = datetime(year=year, month=1, day=1, tzinfo=timezone.utc) return self.get_items(since=ts) def get_items_last90(self) -> Items: @@ -645,10 +651,14 @@ class CollectionItemBase(BaseModel): @property def name(self) -> str: coll = self._collection + if hasattr(coll, "data_type"): + assert coll.data_type name = coll.data_type.value else: + assert coll.merge_type name = coll.merge_type.value + return name def __str__(self): @@ -702,7 +712,9 @@ class CollectionItemBase(BaseModel): @property def path(self) -> FilePath: - return FilePath(os.path.join(self._collection.archive_path, self.filename)) + return_filepath_adapter.validate_python( + os.path.join(self._collection.archive_path, self.filename) + ) @property def partial_path(self) -> FilePath: @@ -732,7 +744,7 @@ class CollectionItemBase(BaseModel): # We assume the target ends with ".####". If not, we'll append .00000 try: - left, right = target.rsplit(".", 1) + _, right = target.rsplit(".", 1) right_int = int(right) except ValueError: return Path(f"{path}.{0:>05}") @@ -740,7 +752,7 @@ class CollectionItemBase(BaseModel): right_int += 1 return Path(f"{path}.{right_int:>05}") - def search_highest_numbered_path(self) -> Optional[Path]: + def search_highest_numbered_path(self) -> Path | None: """This is used for when things are broken, and we want to rebuild our symlinks. We can't trust or use any exist symlinks... so given a path or a partial path... find the highest available "versioned" @@ -766,7 +778,7 @@ class CollectionItemBase(BaseModel): # nums = sorted([b.rsplit(".", 1)[1] for b in builds], reverse=True) # return Path(f"{self.path}.{nums[0]}") - files: List[str] = sorted( + files: list[str] = sorted( builds, key=lambda b: b.rsplit(".", 1)[1], reverse=True ) return Path(os.path.join(coll.archive_path, files[0])) @@ -818,7 +830,6 @@ class CollectionItemBase(BaseModel): shutil.rmtree(generic_path) else: LOG.warning(f"tried removing non-existent file: {generic_path}") - pass def should_archive(self) -> bool: # Determine if enough time has passed to move out of a partial file into an @@ -828,9 +839,7 @@ class CollectionItemBase(BaseModel): if archive_after is None: return False - if datetime.now(tz=timezone.utc) > self.finish + archive_after: - return True - return False + return datetime.now(tz=timezone.utc) > self.finish + archive_after def set_empty(self): assert ( @@ -842,8 +851,8 @@ class CollectionItemBase(BaseModel): def valid_archive( self, - generic_path: Optional[FilePath] = None, - sample: Optional[int] = None, + generic_path: FilePath | None = None, + sample: int | None = None, ) -> bool: """ Attempts to confirm if the parquet file or directory that is @@ -864,7 +873,7 @@ class CollectionItemBase(BaseModel): raise ValueError("Unknown path type.") df = parquet.read().to_pandas() - except (Exception,): + except Exception: LOG.warning(f"Invalid archive {path=}") df = None @@ -876,8 +885,8 @@ class CollectionItemBase(BaseModel): return self.validate_df(df=df, sample=sample) is not None def validate_df( - self, df: pd.DataFrame, sample: Optional[int] = None - ) -> Optional[pd.DataFrame]: + self, df: pd.DataFrame, sample: int | None = None + ) -> pd.DataFrame | None: if sample is not None: sample = min(len(df), sample) try: @@ -910,8 +919,8 @@ class CollectionItemBase(BaseModel): def from_archive( self, include_empty: bool = True, - generic_path: Optional[FilePath] = None, - ) -> Optional[dd.DataFrame]: + generic_path: FilePath | None = None, + ) -> dd.DataFrame | None: if include_empty and self.path_exists(generic_path=self.empty_path): # Return an empty dd.DataFrame with the correct columns @@ -930,15 +939,15 @@ class CollectionItemBase(BaseModel): raise NotImplementedError("Must override") # --- ORM / Data handlers--- - def _to_dict(self, *args, **kwargs) -> dict: - return dict( - should_archive=self.should_archive(), - has_archive=self.has_archive(), - filename=self.filename, - path=self.path, - start=self.start, - finish=self.finish, - ) + def _to_dict(self) -> dict[str, Any]: + return { + "should_archive": self.should_archive(), + "has_archive": self.has_archive(), + "filename": self.filename, + "path": self.path, + "start": self.start, + "finish": self.finish, + } def delete_partial(self): # If a Collection Item is archived, we want to delete the partial file. @@ -961,11 +970,16 @@ class CollectionItemBase(BaseModel): else: self.delete_dangling_partials(keep_latest=2) - def delete_dangling_partials(self, keep_latest=None, target_path=None) -> List[str]: + def delete_dangling_partials( + self, + keep_latest: PositiveInt | None = None, + target_path: Path | str | None = None, + ) -> list[str]: # Specifically looking for numbered partials that are NOT associated # with a symlink. It does not matter if the item is archiveable or not. if target_path is None: target_path = self.partial_path + fps = glob.glob(target_path.as_posix() + ".*") fps = {x for x in fps if x.split(".")[-1].isnumeric()} # Note: if the dir itself is sym-linked, this is going to be wrong. diff --git a/generalresearch/managers/gr/authentication.py b/generalresearch/managers/gr/authentication.py index f4185b2..a402693 100644 --- a/generalresearch/managers/gr/authentication.py +++ b/generalresearch/managers/gr/authentication.py @@ -5,7 +5,6 @@ import logging import os from datetime import datetime, timezone from typing import TYPE_CHECKING, Any -from uuid import uuid4 from psycopg import sql from pydantic import AnyHttpUrl, PositiveInt @@ -23,18 +22,6 @@ if TYPE_CHECKING: class GRUserManager(PostgresManagerWithRedis): - def create_dummy( - self, - sub: str | None = None, - is_superuser: bool = False, - ) -> GRUser: - sub = sub or f"{uuid4().hex}-{uuid4().hex}" - - return self.create( - sub=sub, - is_superuser=is_superuser, - ) - def create( self, sub: str, @@ -71,18 +58,17 @@ class GRUserManager(PostgresManagerWithRedis): def get_by_id(self, gr_user_id: int) -> GRUser | None: from generalresearch.models.gr.authentication import GRUser - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute( - query=""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute( + query=""" SELECT u.* FROM gr_user AS u WHERE u.id = %s LIMIT 1; """, - params=(gr_user_id,), - ) - res = c.fetchone() + params=(gr_user_id,), + ) + res = c.fetchone() if res is None: raise ValueError("GRUser not found") @@ -99,18 +85,17 @@ class GRUserManager(PostgresManagerWithRedis): def get_by_sub(self, sub: str, raises=True) -> GRUser | None: from generalresearch.models.gr.authentication import GRUser - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute( - query=""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute( + query=""" SELECT u.* FROM gr_user AS u WHERE u.sub = %s LIMIT 1; """, - params=(sub,), - ) - res = c.fetchone() + params=(sub,), + ) + res = c.fetchone() if raises and res is None: raise ValueError("GRUser not found") @@ -134,32 +119,30 @@ class GRUserManager(PostgresManagerWithRedis): def get_all(self) -> list[GRUser]: from generalresearch.models.gr.authentication import GRUser - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query=""" + 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() + res = c.fetchall() return [GRUser.from_postgresql(i) for i in res] def get_by_team(self, team_id: PositiveInt) -> list[GRUser]: from generalresearch.models.gr.authentication import GRUser - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute( - query=""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute( + query=""" SELECT gru.* FROM common_membership AS membership INNER JOIN gr_user AS gru ON gru.id = membership.user_id WHERE membership.team_id = %s """, - params=(team_id,), - ) - res = c.fetchall() + params=(team_id,), + ) + res = c.fetchall() for item in res: for k, v in item.items(): @@ -176,7 +159,7 @@ class GRUserManager(PostgresManagerWithRedis): return None res = thl_pg_config.execute_sql_query( - query=f""" + query=""" SELECT bp.id FROM userprofile_brokerageproduct AS bp WHERE bp.business_id = ANY(%s) @@ -240,16 +223,15 @@ class GRTokenManager(PostgresManager): return gr_token # API Key - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - query = sql.SQL(""" + 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 """) - c.execute(query=query, params=(api_key,)) - res = c.fetchall() + c.execute(query=query, params=(api_key,)) + res = c.fetchall() if len(res) == 0: raise Exception(f"No GRUser with token of '{api_key}'") @@ -295,9 +277,8 @@ class GRTokenManager(PostgresManager): # therefore, this will only return 0 or 1 GRTokens from generalresearch.models.gr.authentication import GRToken - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - query = sql.SQL(""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + query = sql.SQL(""" SELECT grt.* FROM gr_token AS grt LEFT JOIN gr_user AS u @@ -306,9 +287,9 @@ class GRTokenManager(PostgresManager): LIMIT 1; """) - c.execute(query=query, params=(user_id,)) + c.execute(query=query, params=(user_id,)) - result = c.fetchall() + result = c.fetchall() if not result: return None diff --git a/generalresearch/managers/gr/business.py b/generalresearch/managers/gr/business.py index aa440fb..4da0e7f 100644 --- a/generalresearch/managers/gr/business.py +++ b/generalresearch/managers/gr/business.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, List +from typing import TYPE_CHECKING from uuid import UUID, uuid4 from psycopg import sql @@ -26,28 +26,6 @@ if TYPE_CHECKING: class BusinessBankAccountManager(PostgresManager): - def create_dummy( - self, - business_id: PositiveInt, - uuid: UUIDStr | None = None, - transfer_method: TransferMethod | None = None, - account_number: str | None = None, - routing_number: str | None = None, - iban: str | None = None, - swift: str | None = None, - ): - from generalresearch.models.gr.business import TransferMethod - - return self.create( - business_id=business_id, - uuid=uuid or uuid4().hex, - transfer_method=transfer_method or TransferMethod.ACH, - account_number=account_number or uuid4().hex[:6], - routing_number=routing_number or uuid4().hex[:6], - iban=iban or uuid4().hex[:6], - swift=swift or uuid4().hex[:6], - ) - def create( self, business_id: PositiveInt, @@ -94,59 +72,25 @@ class BusinessBankAccountManager(PostgresManager): ba.id = ba_id return ba - def get_by_business_id(self, business_id: UUIDStr) -> List[BusinessBankAccount]: + def get_by_business_id(self, business_id: UUIDStr) -> list[BusinessBankAccount]: from generalresearch.models.gr.business import BusinessBankAccount - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute( - query=sql.SQL(""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute( + query=sql.SQL(""" SELECT ba.* FROM common_bankaccount AS ba WHERE ba.business_id = %s """), - params=(business_id,), - ) - res = c.fetchall() + params=(business_id,), + ) + res = c.fetchall() return [BusinessBankAccount.model_validate(item) for item in res] class BusinessAddressManager(PostgresManager): - def create_dummy( - self, - business_id: PositiveInt, - uuid: UUIDStr | None = None, - line_1: str | None = None, - line_2: str | None = None, - city: str | None = None, - state: str | None = None, - postal_code: str | None = None, - phone_number: PhoneNumber | None = None, - country: str | None = None, - ): - uuid = uuid or uuid4().hex - line_1 = line_1 or "abc" - line_2 = line_2 or "bczx" - city = city or "Downingtown" - state = state or "CA" - postal_code = postal_code or "94041" - phone_number = None - country = country or "US" - - return self.create( - business_id=business_id, - uuid=uuid, - line_1=line_1, - line_2=line_2, - city=city, - state=state, - postal_code=postal_code, - phone_number=phone_number, - country=country, - ) - def create( self, business_id: PositiveInt, @@ -237,24 +181,6 @@ class BusinessManager(PostgresManagerWithRedis): uuid=uuid, name=name, team=team, kind=kind, tax_number=tax_number ) - def create_dummy( - self, - uuid: UUIDStr | None = None, - name: str | None = None, - team: Team | None = None, - kind: BusinessType | None = None, - tax_number: str | None = None, - ) -> Business: - from random import randint - - uuid = uuid or uuid4().hex - name = name or "< Unknown >" - tax_number = tax_number or str(randint(1, 999_999_999)) - - return self.create( - uuid=uuid, name=name, team=team, kind=kind, tax_number=tax_number - ) - def create( self, name: str, @@ -316,13 +242,12 @@ class BusinessManager(PostgresManagerWithRedis): """ from generalresearch.models.gr.business import Business - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query=sql.SQL(""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query=sql.SQL(""" SELECT b.id, b.uuid, b.kind, b.name, b.tax_number FROM common_business AS b """)) - res = c.fetchall() + res = c.fetchall() response = [] for i in res: @@ -341,20 +266,19 @@ class BusinessManager(PostgresManagerWithRedis): ) -> list[Business]: # conn: psycopg.Connection = GR_POSTGRES_C.make_connection() - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute( - query=sql.SQL(""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute( + query=sql.SQL(""" SELECT b.id, b.uuid, b.kind, b.name, b.tax_number FROM common_business AS b INNER JOIN common_team_businesses as tb ON tb.business_id = b.id WHERE tb.team_id = %s """), - params=(team_id,), - ) + params=(team_id,), + ) - res = c.fetchall() + res = c.fetchall() response = [] from generalresearch.models.gr.business import Business @@ -372,10 +296,9 @@ class BusinessManager(PostgresManagerWithRedis): ) -> list[Business]: from generalresearch.models.gr.business import Business - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute( - query=sql.SQL(""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute( + query=sql.SQL(""" SELECT b.id, b.uuid, b.kind, b.name, b.tax_number FROM common_business AS b INNER JOIN common_team_businesses AS tb @@ -384,10 +307,10 @@ class BusinessManager(PostgresManagerWithRedis): ON m.team_id = tb.team_id WHERE m.user_id = %s """), - params=(user_id,), - ) + params=(user_id,), + ) - res = c.fetchall() + res = c.fetchall() response = [] for i in res: @@ -402,10 +325,9 @@ class BusinessManager(PostgresManagerWithRedis): :return: Every Business UUIDStr that this GRUser has permission to view """ - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute( - query=sql.SQL(""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute( + query=sql.SQL(""" SELECT b.id FROM common_business AS b INNER JOIN common_team_businesses AS tb @@ -414,10 +336,10 @@ class BusinessManager(PostgresManagerWithRedis): ON tb.team_id = cm.team_id WHERE cm.user_id = %s """), - params=(user_id,), - ) + params=(user_id,), + ) - res = c.fetchall() + res = c.fetchall() return [i["id"] for i in res] @@ -426,10 +348,9 @@ class BusinessManager(PostgresManagerWithRedis): :return: Every Business UUIDStr that this GRUser has permission to view """ - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute( - query=sql.SQL(""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute( + query=sql.SQL(""" SELECT b.uuid FROM common_business AS b INNER JOIN common_team_businesses AS tb @@ -438,10 +359,10 @@ class BusinessManager(PostgresManagerWithRedis): ON tb.team_id = cm.team_id WHERE cm.user_id = %s """), - params=(user_id,), - ) + params=(user_id,), + ) - res = c.fetchall() + res = c.fetchall() return [i["uuid"] for i in res] @@ -453,19 +374,18 @@ class BusinessManager(PostgresManagerWithRedis): assert UUID(hex=business_uuid).hex == business_uuid - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute( - query=sql.SQL(""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute( + query=sql.SQL(""" SELECT id, uuid, kind, name, tax_number FROM common_business WHERE uuid = %s LIMIT 1; """), - params=(business_uuid,), - ) + params=(business_uuid,), + ) - res = c.fetchall() + res = c.fetchall() if len(res) == 0: return None @@ -481,19 +401,18 @@ class BusinessManager(PostgresManagerWithRedis): assert isinstance(business_id, int) - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute( - query=sql.SQL(""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute( + query=sql.SQL(""" SELECT id, uuid, kind, name, tax_number FROM common_business WHERE id = %s LIMIT 1; """), - params=(business_id,), - ) + params=(business_id,), + ) - res = c.fetchall() + res = c.fetchall() if len(res) == 0: return None diff --git a/generalresearch/managers/gr/team.py b/generalresearch/managers/gr/team.py index ecb1ba4..6de82b0 100644 --- a/generalresearch/managers/gr/team.py +++ b/generalresearch/managers/gr/team.py @@ -84,19 +84,18 @@ class MembershipManager(PostgresManager): def exists( self, gr_user_id: PositiveInt, team_id: PositiveInt ) -> Membership | None: - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute( - query=sql.SQL(""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute( + query=sql.SQL(""" SELECT id, uuid, privilege, owner, created, user_id, team_id FROM common_membership WHERE team_id = %s AND user_id = %s LIMIT 1 """), - params=(team_id, gr_user_id), - ) - res = c.fetchone() + params=(team_id, gr_user_id), + ) + res = c.fetchone() if not res: return None @@ -104,36 +103,34 @@ class MembershipManager(PostgresManager): return Membership.model_validate(res) def get_by_team_id(self, team_id: PositiveInt) -> list[Membership]: - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute( - query=sql.SQL(""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute( + query=sql.SQL(""" SELECT id, uuid, privilege, owner, created, user_id, team_id FROM common_membership WHERE team_id = %s LIMIT 250 """), - params=(team_id,), - ) - res = c.fetchall() + params=(team_id,), + ) + res = c.fetchall() return [Membership.model_validate(i) for i in res] def get_by_gr_user_id(self, gr_user_id: PositiveInt) -> list[Membership]: - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute( - query=sql.SQL(""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute( + query=sql.SQL(""" SELECT id, uuid, privilege, owner, created, user_id, team_id FROM common_membership WHERE user_id = %s LIMIT 250 """), - params=(gr_user_id,), - ) - res = c.fetchall() + params=(gr_user_id,), + ) + res = c.fetchall() return [Membership.model_validate(i) for i in res] @@ -154,24 +151,15 @@ class TeamManager(PostgresManagerWithRedis): def get_all(self) -> list[Team]: from generalresearch.models.gr.team import Team - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute(query=sql.SQL(""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute(query=sql.SQL(""" SELECT t.id, t.uuid, t.name FROM common_team AS t """)) - res = c.fetchall() + res = c.fetchall() return [Team.model_validate(i) for i in res] - def create_dummy( - self, uuid: UUIDStr | None = None, name: str | None = None - ) -> Team: - uuid = uuid or uuid4().hex - name = name or f"name-{uuid4().hex[:12]}" - - return self.create(uuid=uuid, name=name) - def create( self, name: str, @@ -228,19 +216,18 @@ class TeamManager(PostgresManagerWithRedis): def get_by_uuid(self, team_uuid: UUIDStr) -> Team | None: from generalresearch.models.gr.team import Team - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute( - query=""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute( + query=""" SELECT t.* FROM common_team AS t WHERE t.uuid = %s LIMIT 1; """, - params=(team_uuid,), - ) + params=(team_uuid,), + ) - res = c.fetchone() + res = c.fetchone() if not isinstance(res, dict): return None @@ -250,19 +237,18 @@ class TeamManager(PostgresManagerWithRedis): def get_by_id(self, team_id: PositiveInt) -> Team | None: from generalresearch.models.gr.team import Team - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute( - query=sql.SQL(""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute( + query=sql.SQL(""" SELECT t.id, t.uuid, t.name FROM common_team AS t WHERE t.id = %s LIMIT 1; """), - params=(team_id,), - ) + params=(team_id,), + ) - res = c.fetchone() + res = c.fetchone() if not isinstance(res, dict): return None @@ -272,19 +258,18 @@ class TeamManager(PostgresManagerWithRedis): def get_by_user(self, gr_user: GRUser) -> list[Team]: from generalresearch.models.gr.team import Team - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - c.execute( - query=sql.SQL(""" + with self.pg_config.make_connection() as conn, conn.cursor() as c: + c.execute( + query=sql.SQL(""" SELECT team.* FROM common_team AS team INNER JOIN common_membership AS mem ON mem.team_id = team.id WHERE mem.user_id = %s """), - params=(gr_user.id,), - ) + params=(gr_user.id,), + ) - res = c.fetchall() + res = c.fetchall() return [Team.model_validate(item) for item in res] diff --git a/generalresearch/managers/leaderboard/__init__.py b/generalresearch/managers/leaderboard/__init__.py index aae4a05..d5138cd 100644 --- a/generalresearch/managers/leaderboard/__init__.py +++ b/generalresearch/managers/leaderboard/__init__.py @@ -1,4 +1,4 @@ -from typing import Dict +from __future__ import annotations import pytz from cachetools import LRUCache, cached @@ -6,7 +6,7 @@ from zoneinfo import ZoneInfo @cached(cache=LRUCache(maxsize=1)) -def country_timezone() -> Dict[str, ZoneInfo]: +def country_timezone() -> dict[str, ZoneInfo]: """ Most countries only have 1 tz. I am picking the most populous for the rest. A timezone is unique for a country, as in America/New_York and America/Toronto diff --git a/generalresearch/managers/thl/ipinfo.py b/generalresearch/managers/thl/ipinfo.py index 510dc63..e1143c2 100644 --- a/generalresearch/managers/thl/ipinfo.py +++ b/generalresearch/managers/thl/ipinfo.py @@ -3,7 +3,6 @@ from __future__ import annotations import ipaddress from collections.abc import Collection from decimal import Decimal -from random import randint import faker import pymysql @@ -33,38 +32,6 @@ fake = faker.Faker() class IPGeonameManager(PostgresManager): - def create_dummy( - self, - geoname_id: PositiveInt | None = None, - continent_code: str | None = None, - continent_name: str | None = None, - country_iso: str | None = None, - country_name: str | None = None, - subdivision_1_iso: str | None = None, - subdivision_1_name: str | None = None, - subdivision_2_iso: str | None = None, - subdivision_2_name: str | None = None, - city_name: str | None = None, - metro_code: int | None = None, - time_zone: str | None = None, - is_in_european_union: bool | None = None, - ) -> IPGeoname: - return self.create( - geoname_id=geoname_id or randint(1, 999_999_999), - continent_code=continent_code or "na", - continent_name=continent_name or "North America", - country_iso=country_iso or "us", - country_name=country_name or "United States", - subdivision_1_iso=subdivision_1_iso or "fl", - subdivision_1_name=subdivision_1_name or "Florida", - subdivision_2_iso=subdivision_2_iso, - subdivision_2_name=subdivision_2_name, - city_name=city_name, - metro_code=metro_code, - time_zone=time_zone, - is_in_european_union=is_in_european_union, - ) - def create_basic( self, geoname_id: PositiveInt, @@ -230,60 +197,6 @@ class IPGeonameManager(PostgresManager): class IPInformationManager(PostgresManager): - def create_dummy( - self, - ip: IPvAnyAddressStr | None = None, - geoname_id: PositiveInt | None = None, - country_iso: str | None = None, - registered_country_iso: str | None = None, - is_anonymous: bool | None = None, - is_anonymous_vpn: bool | None = None, - is_hosting_provider: bool | None = None, - is_public_proxy: bool | None = None, - is_tor_exit_node: bool | None = None, - is_residential_proxy: bool | None = None, - autonomous_system_number: PositiveInt | None = None, - autonomous_system_organization: str | None = None, - domain: str | None = None, - isp: str | None = None, - mobile_country_code: str | None = None, - mobile_network_code: str | None = None, - network: str | None = None, - organization: str | None = None, - static_ip_score: float | None = None, - user_type: UserType | None = None, - postal_code: str | None = None, - latitude: Decimal | None = None, - longitude: Decimal | None = None, - accuracy_radius: int | None = None, - ) -> IPInformation: - return self.create( - ip=ip or fake.ipv4_public(), - geoname_id=geoname_id, - country_iso=country_iso or fake.country_code(), - registered_country_iso=registered_country_iso, - is_anonymous=is_anonymous, - is_anonymous_vpn=is_anonymous_vpn, - is_hosting_provider=is_hosting_provider, - is_public_proxy=is_public_proxy, - is_tor_exit_node=is_tor_exit_node, - is_residential_proxy=is_residential_proxy, - autonomous_system_number=autonomous_system_number, - autonomous_system_organization=autonomous_system_organization, - domain=domain, - isp=isp, - mobile_country_code=mobile_country_code, - mobile_network_code=mobile_network_code, - network=network, - organization=organization, - static_ip_score=static_ip_score, - user_type=user_type, - postal_code=postal_code, - latitude=latitude, - longitude=longitude, - accuracy_radius=accuracy_radius, - ) - def create_basic( self, ip: IPvAnyAddressStr, @@ -440,16 +353,15 @@ class IPInformationManager(PostgresManager): if len(filter_ips) == 0: return [] - with self.pg_config.make_connection() as conn: - with conn.cursor() as c: - res = [] - for chunk in chunked(filter_ips, 500): - res.extend( - self.fetch_ip_information_( - c=c, - filter_ips=chunk, - ) + with self.pg_config.make_connection() as conn, conn.cursor() as c: + res = [] + for chunk in chunked(filter_ips, 500): + res.extend( + self.fetch_ip_information_( + c=c, + filter_ips=chunk, ) + ) return res def fetch_ip_information_( @@ -518,8 +430,6 @@ class IPInformationManager(PostgresManager): # percent_empty = numerator / (denominator or 1) # TODO: Post to telegraf / grafana - return None - class GeoIpInfoManager(PostgresManagerWithRedis): diff --git a/pyproject.toml b/pyproject.toml index dd2e649..183e271 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,7 @@ description = "Python Utilities for General Research" readme = "README.md" requires-python = ">=3.8" dependencies = [ + "fastapi", "Faker", "PyMySQL", "psycopg", diff --git a/test_utils/managers/gr/__init__.py b/test_utils/managers/gr/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/test_utils/managers/gr/__init__.py diff --git a/test_utils/managers/gr/conftest.py b/test_utils/managers/gr/conftest.py new file mode 100644 index 0000000..3e3a4ad --- /dev/null +++ b/test_utils/managers/gr/conftest.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +from typing import Callable +from uuid import uuid4 + +import pytest +from pydantic import PositiveInt +from pydantic_extra_types.phone_numbers import PhoneNumber + +from generalresearch.managers.gr.authentication import GRUserManager +from generalresearch.managers.gr.business import ( + BusinessAddressManager, + BusinessBankAccountManager, + BusinessManager, +) +from generalresearch.managers.gr.team import TeamManager +from generalresearch.models.custom_types import UUIDStr +from generalresearch.models.gr.authentication import GRUser +from generalresearch.models.gr.business import ( + Business, + BusinessAddress, + BusinessBankAccount, + BusinessType, + TransferMethod, +) +from generalresearch.models.gr.team import Team + + +@pytest.fixture +def gr_user_factory(gr_um: GRUserManager) -> Callable[..., GRUser]: + + def _inner( + sub: str | None = None, + is_superuser: bool = False, + ) -> GRUser: + sub = sub or f"{uuid4().hex}-{uuid4().hex}" + + return gr_um.create( + sub=sub, + is_superuser=is_superuser, + ) + + return _inner + + +@pytest.fixture +def gr_business_bank_account_factory( + gr_bbam: BusinessBankAccountManager, +) -> Callable[..., BusinessBankAccount]: + + def _inner( + business_id: PositiveInt, + uuid: UUIDStr | None = None, + transfer_method: TransferMethod | None = None, + account_number: str | None = None, + routing_number: str | None = None, + iban: str | None = None, + swift: str | None = None, + ): + from generalresearch.models.gr.business import TransferMethod + + return gr_bbam.create( + business_id=business_id, + uuid=uuid or uuid4().hex, + transfer_method=transfer_method or TransferMethod.ACH, + account_number=account_number or uuid4().hex[:6], + routing_number=routing_number or uuid4().hex[:6], + iban=iban or uuid4().hex[:6], + swift=swift or uuid4().hex[:6], + ) + + return _inner + + +@pytest.fixture +def gr_business_address_factory( + gr_bam: BusinessAddressManager, +) -> Callable[..., BusinessAddress]: + + def _inner( + business_id: PositiveInt, + uuid: UUIDStr | None = None, + line_1: str | None = None, + line_2: str | None = None, + city: str | None = None, + state: str | None = None, + postal_code: str | None = None, + phone_number: PhoneNumber | None = None, + country: str | None = None, + ): + uuid = uuid or uuid4().hex + line_1 = line_1 or "abc" + line_2 = line_2 or "bczx" + city = city or "Downingtown" + state = state or "CA" + postal_code = postal_code or "94041" + phone_number = None + country = country or "US" + + return gr_bam.create( + business_id=business_id, + uuid=uuid, + line_1=line_1, + line_2=line_2, + city=city, + state=state, + postal_code=postal_code, + phone_number=phone_number, + country=country, + ) + + return _inner + + +@pytest.fixture +def gr_business_factory( + gr_bm: BusinessManager, +) -> Callable[..., Business]: + + def _inner( + uuid: UUIDStr | None = None, + name: str | None = None, + team: Team | None = None, + kind: BusinessType | None = None, + tax_number: str | None = None, + ) -> Business: + from random import randint + + uuid = uuid or uuid4().hex + name = name or "< Unknown >" + tax_number = tax_number or str(randint(1, 999_999_999)) + + return gr_bm.create( + uuid=uuid, name=name, team=team, kind=kind, tax_number=tax_number + ) + + return _inner + + +@pytest.fixture +def gr_team( + gr_tm: TeamManager, +) -> Callable[..., Team]: + + def _inner(uuid: UUIDStr | None = None, name: str | None = None) -> Team: + uuid = uuid or uuid4().hex + name = name or f"name-{uuid4().hex[:12]}" + + return gr_tm.create(uuid=uuid, name=name) + + return _inner diff --git a/test_utils/managers/grliq/__init__.py b/test_utils/managers/grliq/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/test_utils/managers/grliq/__init__.py diff --git a/test_utils/managers/grliq/conftest.py b/test_utils/managers/grliq/conftest.py new file mode 100644 index 0000000..525d8c8 --- /dev/null +++ b/test_utils/managers/grliq/conftest.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Callable +from uuid import uuid4 + +import pytest + +from generalresearch.grliq.managers import DUMMY_GRLIQ_DATA +from generalresearch.grliq.managers.forensic_data import GrlIqDataManager +from generalresearch.grliq.models.forensic_data import GrlIqData + + +@pytest.fixture +def grliq_data_factory(grliq_dm: GrlIqDataManager) -> Callable[..., GrlIqData]: + + def _inner( + is_attempt_allowed: bool = True, + product_id: str | None = None, + product_user_id: str | None = None, + uuid: str | None = None, + mid: str | None = None, + created_at: datetime | None = None, + ) -> GrlIqData: + """ + Creates a dummy record in the db with a GrlIqData (data), GrlIqCheckerResults (result_data), + and GrlIqForensicCategoryResult (category_results) + :param is_attempt_allowed: Whether the attempt is allowed. + :param product_id: product_id of user + :param product_user_id: product_user_id of user + :param uuid: uuid for the grliq data record + :param mid: the thl_session:uuid / mid for the attempt. + :return: + """ + import copy + + res: GrlIqData = copy.deepcopy(DUMMY_GRLIQ_DATA[int(is_attempt_allowed)]) + + product_id = product_id or uuid4().hex + product_user_id = product_user_id or uuid4().hex + uuid = uuid or uuid4().hex + mid = mid or uuid4().hex + created_at = created_at or datetime.now(tz=timezone.utc) + + res["data"].product_id = product_id + res["data"].product_user_id = product_user_id + res["data"].uuid = uuid + res["data"].mid = mid + res["data"].created_at = created_at + res["result_data"].uuid = uuid + res["category_result"].uuid = uuid + + return grliq_dm.create( + iq_data=res["data"], + result_data=res["result_data"], + category_result=res["category_result"], + fraud_score=res["category_result"].fraud_score, + is_attempt_allowed=res["category_result"].is_attempt_allowed(), + ) + + return _inner diff --git a/test_utils/managers/thl/__init__.py b/test_utils/managers/thl/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/test_utils/managers/thl/__init__.py diff --git a/test_utils/managers/thl/conftest.py b/test_utils/managers/thl/conftest.py new file mode 100644 index 0000000..76e4226 --- /dev/null +++ b/test_utils/managers/thl/conftest.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +from decimal import Decimal +from random import randint +from typing import Callable + +import faker +from pydantic import PositiveInt + +from generalresearch.managers.thl.ipinfo import IPGeonameManager, IPInformationManager +from generalresearch.models.custom_types import IPvAnyAddressStr +from generalresearch.models.thl.ipinfo import IPGeoname, IPInformation, UserType + +fake = faker.Faker() + + +def ipgeoname_factory(ipgeoname_manager: IPGeonameManager) -> Callable[..., IPGeoname]: + + def _inner( + geoname_id: PositiveInt | None = None, + continent_code: str | None = None, + continent_name: str | None = None, + country_iso: str | None = None, + country_name: str | None = None, + subdivision_1_iso: str | None = None, + subdivision_1_name: str | None = None, + subdivision_2_iso: str | None = None, + subdivision_2_name: str | None = None, + city_name: str | None = None, + metro_code: int | None = None, + time_zone: str | None = None, + is_in_european_union: bool | None = None, + ) -> IPGeoname: + + return ipgeoname_manager.create( + geoname_id=geoname_id or randint(1, 999_999_999), + continent_code=continent_code or "na", + continent_name=continent_name or "North America", + country_iso=country_iso or "us", + country_name=country_name or "United States", + subdivision_1_iso=subdivision_1_iso or "fl", + subdivision_1_name=subdivision_1_name or "Florida", + subdivision_2_iso=subdivision_2_iso, + subdivision_2_name=subdivision_2_name, + city_name=city_name, + metro_code=metro_code, + time_zone=time_zone, + is_in_european_union=is_in_european_union, + ) + + return _inner + + +def ipinformation_factory( + ipinformation_manager: IPInformationManager, +) -> Callable[..., IPInformation]: + + def _inner( + ip: IPvAnyAddressStr | None = None, + geoname_id: PositiveInt | None = None, + country_iso: str | None = None, + registered_country_iso: str | None = None, + is_anonymous: bool | None = None, + is_anonymous_vpn: bool | None = None, + is_hosting_provider: bool | None = None, + is_public_proxy: bool | None = None, + is_tor_exit_node: bool | None = None, + is_residential_proxy: bool | None = None, + autonomous_system_number: PositiveInt | None = None, + autonomous_system_organization: str | None = None, + domain: str | None = None, + isp: str | None = None, + mobile_country_code: str | None = None, + mobile_network_code: str | None = None, + network: str | None = None, + organization: str | None = None, + static_ip_score: float | None = None, + user_type: UserType | None = None, + postal_code: str | None = None, + latitude: Decimal | None = None, + longitude: Decimal | None = None, + accuracy_radius: int | None = None, + ) -> IPInformation: + + return ipinformation_manager.create( + ip=ip or fake.ipv4_public(), + geoname_id=geoname_id, + country_iso=country_iso or fake.country_code(), + registered_country_iso=registered_country_iso, + is_anonymous=is_anonymous, + is_anonymous_vpn=is_anonymous_vpn, + is_hosting_provider=is_hosting_provider, + is_public_proxy=is_public_proxy, + is_tor_exit_node=is_tor_exit_node, + is_residential_proxy=is_residential_proxy, + autonomous_system_number=autonomous_system_number, + autonomous_system_organization=autonomous_system_organization, + domain=domain, + isp=isp, + mobile_country_code=mobile_country_code, + mobile_network_code=mobile_network_code, + network=network, + organization=organization, + static_ip_score=static_ip_score, + user_type=user_type, + postal_code=postal_code, + latitude=latitude, + longitude=longitude, + accuracy_radius=accuracy_radius, + ) + + return _inner diff --git a/test_utils/models/conftest.py b/test_utils/models/conftest.py index 64bdec6..1133d32 100644 --- a/test_utils/models/conftest.py +++ b/test_utils/models/conftest.py @@ -1,11 +1,14 @@ +from __future__ import annotations + from datetime import datetime, timedelta, timezone from decimal import Decimal from random import choice as randchoice from random import randint -from typing import TYPE_CHECKING, Callable, Dict, List, Optional +from typing import TYPE_CHECKING, Callable from uuid import uuid4 import pytest +from fastapi import Request from pydantic import AwareDatetime, PositiveInt from generalresearch.models import Source @@ -84,26 +87,27 @@ def user( @pytest.fixture def user_with_wallet( - request, user_factory: Callable[..., "User"], product_user_wallet_yes: "Product" -) -> "User": + user_factory: Callable[..., User], + product_user_wallet_yes: Product, +) -> User: # A user on a product with user wallet enabled, but they have no money return user_factory(product=product_user_wallet_yes) @pytest.fixture def user_with_wallet_amt( - request, user_factory: Callable[..., "User"], product_amt_true: "Product" -) -> "User": + user_factory: Callable[..., User], product_amt_true: Product +) -> User: # A user on a product with user wallet enabled, on AMT, but they have no money return user_factory(product=product_amt_true) @pytest.fixture(scope="function") def user_factory( - user_manager: "UserManager", thl_web_rr: PostgresConfig -) -> Callable[..., "User"]: + user_manager: UserManager, thl_web_rr: PostgresConfig +) -> Callable[..., User]: - def _inner(product: "Product", created: Optional[datetime] = None) -> "User": + def _inner(product: Product, created: datetime | None = None) -> User: u = user_manager.create_dummy(product=product, created=created) u.prefetch_product(pg_config=thl_web_rr) @@ -113,11 +117,11 @@ def user_factory( @pytest.fixture -def wall_factory(wall_manager: "WallManager") -> Callable[..., "Wall"]: +def wall_factory(wall_manager: WallManager) -> Callable[..., Wall]: def _inner( - session: "Session", wall_status: "Status", req_cpi: Optional[Decimal] = None - ) -> "Wall": + session: Session, wall_status: Status, req_cpi: Decimal | None = None + ) -> Wall: assert session.started <= datetime.now( tz=timezone.utc @@ -153,9 +157,7 @@ def wall_factory(wall_manager: "WallManager") -> Callable[..., "Wall"]: @pytest.fixture -def wall( - session: "Session", user: "User", wall_manager: "WallManager" -) -> Optional["Wall"]: +def wall(session: Session, user: User, wall_manager: WallManager) -> Wall | None: from generalresearch.models.thl.task_status import StatusCode1 wall = wall_manager.create_dummy(session_id=session.id, user_id=user.user_id) @@ -170,20 +172,20 @@ def wall( @pytest.fixture def session_factory( - wall_factory: Callable[..., "Wall"], - session_manager: "SessionManager", - wall_manager: "WallManager", + wall_factory: Callable[..., Wall], + session_manager: SessionManager, + wall_manager: WallManager, utc_hour_ago: datetime, -) -> Callable[..., "Session"]: +) -> Callable[..., Session]: from generalresearch.models.thl.session import Source def _inner( - user: "User", + user: User, # Wall details wall_count: int = 5, wall_req_cpi: Decimal = Decimal(".50"), - wall_req_cpis: Optional[List[Decimal]] = None, - wall_statuses: Optional[List[Status]] = None, + wall_req_cpis: list[Decimal] | None = None, + wall_statuses: list[Status] | None = None, wall_source: Source = Source.TESTING, # Session details final_status: Status = Status.COMPLETE, @@ -236,24 +238,24 @@ def session_factory( @pytest.fixture(scope="function") def finished_session_factory( - session_factory: Callable[..., "Session"], - session_manager: "SessionManager", + session_factory: Callable[..., Session], + session_manager: SessionManager, utc_hour_ago: datetime, -) -> Callable[..., "Session"]: +) -> Callable[..., Session]: from generalresearch.models.thl.session import Source def _inner( - user: "User", + user: User, # Wall details wall_count: int = 5, wall_req_cpi: Decimal = Decimal(".50"), - wall_req_cpis: Optional[List[Decimal]] = None, - wall_statuses: Optional[List[Status]] = None, + wall_req_cpis: list[Decimal] | None = None, + wall_statuses: list[Status] | None = None, wall_source: Source = Source.TESTING, # Session details final_status: Status = Status.COMPLETE, started: datetime = utc_hour_ago, - ) -> "Session": + ) -> Session: s: Session = session_factory( user=user, wall_count=wall_count, @@ -281,9 +283,8 @@ def finished_session_factory( @pytest.fixture def session( - user: "User", session_manager: "SessionManager", wall_manager: "WallManager" -) -> "Session": - from generalresearch.models.thl.session import Session, Wall + user: User, session_manager: SessionManager, wall_manager: WallManager +) -> Session: session: Session = session_manager.create_dummy(user=user, country_iso="us") wall: Wall = wall_manager.create_dummy( @@ -297,7 +298,7 @@ def session( @pytest.fixture -def product(request, product_manager: "ProductManager") -> "Product": +def product(request: Request, product_manager: ProductManager) -> Product: team = getattr(request, "team", None) business = getattr(request, "business", None) @@ -309,13 +310,13 @@ def product(request, product_manager: "ProductManager") -> "Product": @pytest.fixture -def product_factory(product_manager: "ProductManager") -> Callable[..., "Product"]: +def product_factory(product_manager: ProductManager) -> Callable[..., Product]: def _inner( - team: Optional["Team"] = None, - business: Optional["Business"] = None, + team: Team | None = None, + business: Business | None = None, commission_pct: Decimal = Decimal("0.05"), - ) -> "Product": + ) -> Product: return product_manager.create_dummy( team_id=team.uuid if team else None, business_id=business.uuid if business else None, @@ -326,7 +327,7 @@ def product_factory(product_manager: "ProductManager") -> Callable[..., "Product @pytest.fixture -def payout_config(request) -> "PayoutConfig": +def payout_config(request: Request) -> PayoutConfig: from generalresearch.models.thl.product import ( PayoutConfig, PayoutTransformation, @@ -348,8 +349,8 @@ def payout_config(request) -> "PayoutConfig": @pytest.fixture def product_user_wallet_yes( - payout_config: "PayoutConfig", product_manager: "ProductManager" -) -> "Product": + payout_config: PayoutConfig, product_manager: ProductManager +) -> Product: from generalresearch.models.thl.product import UserWalletConfig return product_manager.create_dummy( @@ -358,7 +359,7 @@ def product_user_wallet_yes( @pytest.fixture -def product_user_wallet_no(product_manager: "ProductManager") -> "Product": +def product_user_wallet_no(product_manager: ProductManager) -> Product: from generalresearch.models.thl.product import UserWalletConfig return product_manager.create_dummy( @@ -368,8 +369,8 @@ def product_user_wallet_no(product_manager: "ProductManager") -> "Product": @pytest.fixture def product_amt_true( - product_manager: "ProductManager", payout_config: "PayoutConfig" -) -> "Product": + product_manager: ProductManager, payout_config: PayoutConfig +) -> Product: from generalresearch.models.thl.product import UserWalletConfig return product_manager.create_dummy( @@ -380,19 +381,19 @@ def product_amt_true( @pytest.fixture def bp_payout_factory( - thl_lm: "ThlLedgerManager", - product_manager: "ProductManager", - business_payout_event_manager: "BusinessPayoutEventManager", -) -> Callable[..., "BrokerageProductPayoutEvent"]: + thl_lm: ThlLedgerManager, + product_manager: ProductManager, + business_payout_event_manager: BusinessPayoutEventManager, +) -> Callable[..., BrokerageProductPayoutEvent]: def _inner( - product: Optional["Product"] = None, - amount: Optional["USDCent"] = None, - ext_ref_id: Optional[str] = None, - created: Optional[AwareDatetime] = None, + product: Product | None = None, + amount: USDCent | None = None, + ext_ref_id: str | None = None, + created: AwareDatetime | None = None, skip_wallet_balance_check: bool = False, skip_one_per_day_check: bool = False, - ) -> "BrokerageProductPayoutEvent": + ) -> BrokerageProductPayoutEvent: from generalresearch.currency import USDCent product = product or product_manager.create_dummy() @@ -415,43 +416,43 @@ def bp_payout_factory( @pytest.fixture -def business(request, business_manager: "BusinessManager") -> "Business": +def business(request, business_manager: BusinessManager) -> Business: return business_manager.create_dummy() @pytest.fixture def business_address( - request, business: "Business", business_address_manager: "BusinessAddressManager" -) -> "BusinessAddress": + request, business: "Business", business_address_manager: BusinessAddressManager +) -> BusinessAddress: return business_address_manager.create_dummy(business_id=business.id) @pytest.fixture def business_bank_account( request, - business: "Business", - business_bank_account_manager: "BusinessBankAccountManager", -) -> "BusinessBankAccount": + business: Business, + business_bank_account_manager: BusinessBankAccountManager, +) -> BusinessBankAccount: return business_bank_account_manager.create_dummy(business_id=business.id) @pytest.fixture -def team(request, team_manager: "TeamManager") -> "Team": +def team(request, team_manager: TeamManager) -> Team: return team_manager.create_dummy() @pytest.fixture -def gr_user(gr_um: "GRUserManager") -> "GRUser": +def gr_user(gr_um: GRUserManager) -> GRUser: return gr_um.create_dummy() @pytest.fixture def gr_user_cache( - gr_user: "GRUser", + gr_user: GRUser, gr_db: PostgresConfig, thl_web_rr: PostgresConfig, gr_redis_config: RedisConfig, -) -> "GRUser": +) -> GRUser: gr_user.set_cache( pg_config=gr_db, thl_web_rr=thl_web_rr, redis_config=gr_redis_config ) @@ -459,7 +460,7 @@ def gr_user_cache( @pytest.fixture -def gr_user_factory(gr_um: "GRUserManager") -> Callable[..., "GRUser"]: +def gr_user_factory(gr_um: GRUserManager) -> Callable[..., GRUser]: def _inner(): return gr_um.create_dummy() @@ -469,8 +470,8 @@ def gr_user_factory(gr_um: "GRUserManager") -> Callable[..., "GRUser"]: @pytest.fixture() def gr_user_token( - gr_user: "GRUser", gr_tm: "GRTokenManager", gr_db: PostgresConfig -) -> "GRToken": + gr_user: GRUser, gr_tm: GRTokenManager, gr_db: PostgresConfig +) -> GRToken: gr_tm.create(user_id=gr_user.id) gr_user.prefetch_token(pg_config=gr_db) @@ -480,14 +481,14 @@ def gr_user_token( @pytest.fixture() -def gr_user_token_header(gr_user_token: "GRToken") -> Dict[str, str]: +def gr_user_token_header(gr_user_token: GRToken) -> dict[str, str]: return gr_user_token.auth_header @pytest.fixture(scope="function") def membership( - request, team: "Team", gr_user: "GRUser", team_manager: "TeamManager" -) -> "Membership": + request, team: Team, gr_user: GRUser, team_manager: TeamManager +) -> Membership: assert team.id, "Team must be saved" assert gr_user.id, "GRUser must be saved" return team_manager.add_user(team=team, gr_user=gr_user) @@ -495,14 +496,14 @@ def membership( @pytest.fixture(scope="function") def membership_factory( - team: "Team", - gr_user: "GRUser", - membership_manager: "MembershipManager", - team_manager: "TeamManager", - gr_um: "GRUserManager", -) -> Callable[..., "Membership"]: - - def _inner(**kwargs) -> "Membership": + team: Team, + gr_user: GRUser, + membership_manager: MembershipManager, + team_manager: TeamManager, + gr_um: GRUserManager, +) -> Callable[..., Membership]: + + def _inner(**kwargs) -> Membership: _team = kwargs.get("team", team_manager.create_dummy()) _gr_user = kwargs.get("gr_user", gr_um.create_dummy()) @@ -512,23 +513,23 @@ def membership_factory( @pytest.fixture -def audit_log(audit_log_manager: "AuditLogManager", user: "User") -> "AuditLog": +def audit_log(audit_log_manager: AuditLogManager, user: User) -> AuditLog: return audit_log_manager.create_dummy(user_id=user.user_id) @pytest.fixture def audit_log_factory( - audit_log_manager: "AuditLogManager", -) -> Callable[..., "AuditLog"]: + audit_log_manager: AuditLogManager, +) -> Callable[..., AuditLog]: def _inner( user_id: PositiveInt, - level: Optional["AuditLogLevel"] = None, - event_type: Optional[str] = None, - event_msg: Optional[str] = None, - event_value: Optional[float] = None, - ) -> "AuditLog": + level: AuditLogLevel | None = None, + event_type: str | None = None, + event_msg: str | None = None, + event_value: float | None = None, + ) -> AuditLog: return audit_log_manager.create_dummy( user_id=user_id, level=level, @@ -541,14 +542,14 @@ def audit_log_factory( @pytest.fixture -def ip_geoname(ip_geoname_manager: "IPGeonameManager") -> "IPGeoname": +def ip_geoname(ip_geoname_manager: IPGeonameManager) -> IPGeoname: return ip_geoname_manager.create_dummy() @pytest.fixture def ip_information( - ip_information_manager: "IPInformationManager", ip_geoname: "IPGeoname" -) -> "IPInformation": + ip_information_manager: IPInformationManager, ip_geoname: IPGeoname +) -> IPInformation: return ip_information_manager.create_dummy( geoname_id=ip_geoname.geoname_id, country_iso=ip_geoname.country_iso ) @@ -556,10 +557,10 @@ def ip_information( @pytest.fixture def ip_information_factory( - ip_information_manager: "IPInformationManager", -) -> Callable[..., "IPInformation"]: + ip_information_manager: IPInformationManager, +) -> Callable[..., IPInformation]: - def _inner(ip: str, geoname: "IPGeoname", **kwargs) -> "IPInformation": + def _inner(ip: str, geoname: IPGeoname, **kwargs) -> IPInformation: return ip_information_manager.create_dummy( ip=ip, geoname_id=geoname.geoname_id, @@ -572,25 +573,25 @@ def ip_information_factory( @pytest.fixture def ip_record( - ip_record_manager: "IPRecordManager", ip_geoname: "IPGeoname", user: "User" -) -> "IPRecord": + ip_record_manager: IPRecordManager, ip_geoname: IPGeoname, user: User +) -> IPRecord: return ip_record_manager.create_dummy(user_id=user.user_id) @pytest.fixture def ip_record_factory( - ip_record_manager: "IPRecordManager", user: "User" -) -> Callable[..., "IPRecord"]: + ip_record_manager: IPRecordManager, user: User +) -> Callable[..., IPRecord]: - def _inner(user_id: PositiveInt, ip: Optional[str] = None) -> "IPRecord": + def _inner(user_id: PositiveInt, ip: str | None = None) -> IPRecord: return ip_record_manager.create_dummy(user_id=user_id, ip=ip) return _inner @pytest.fixture(scope="session") -def buyer(buyer_manager: "BuyerManager") -> "Buyer": +def buyer(buyer_manager: BuyerManager) -> Buyer: buyer_code = uuid4().hex buyer_manager.bulk_get_or_create(source=Source.TESTING, codes=[buyer_code]) b = Buyer( @@ -601,7 +602,7 @@ def buyer(buyer_manager: "BuyerManager") -> "Buyer": @pytest.fixture(scope="session") -def buyer_factory(buyer_manager: "BuyerManager") -> Callable[..., "Buyer"]: +def buyer_factory(buyer_manager: BuyerManager) -> Callable[..., Buyer]: def _inner() -> Buyer: return buyer_manager.bulk_get_or_create( @@ -612,7 +613,7 @@ def buyer_factory(buyer_manager: "BuyerManager") -> Callable[..., "Buyer"]: @pytest.fixture(scope="session") -def survey(survey_manager: "SurveyManager", buyer: "Buyer") -> "Survey": +def survey(survey_manager: SurveyManager, buyer: Buyer) -> "Survey": s = Survey(source=Source.TESTING, survey_id=uuid4().hex, buyer_code=buyer.code) survey_manager.create_bulk([s]) return s @@ -620,10 +621,10 @@ def survey(survey_manager: "SurveyManager", buyer: "Buyer") -> "Survey": @pytest.fixture(scope="session") def survey_factory( - survey_manager: "SurveyManager", buyer_factory: Callable[..., "Buyer"] -) -> Callable[..., "Survey"]: + survey_manager: SurveyManager, buyer_factory: Callable[..., Buyer] +) -> Callable[..., Survey]: - def _inner(buyer: Optional[Buyer] = None) -> "Survey": + def _inner(buyer: Buyer | None = None) -> Survey: buyer = buyer or buyer_factory() s = Survey( source=Source.TESTING, diff --git a/tests/incite/test_collection_base.py b/tests/incite/test_collection_base.py index 497e5ab..7e6605f 100644 --- a/tests/incite/test_collection_base.py +++ b/tests/incite/test_collection_base.py @@ -1,5 +1,6 @@ -from datetime import datetime, timezone, timedelta -from os.path import exists as pexists, join as pjoin +from datetime import datetime, timedelta, timezone +from os.path import exists as pexists +from os.path import join as pjoin from pathlib import Path from uuid import uuid4 @@ -244,6 +245,7 @@ class TestCollectionBaseMethodsCleanup: class TestCollectionBaseMethodsCleanup: + @pytest.mark.skip def test_cleanup_partials(self, mnt_filepath): instance = CollectionBase(archive_path=mnt_filepath.data_src) diff --git a/tests/models/custom_types/test_aware_datetime.py b/tests/models/custom_types/test_aware_datetime.py index 14d1343..530142e 100644 --- a/tests/models/custom_types/test_aware_datetime.py +++ b/tests/models/custom_types/test_aware_datetime.py @@ -1,6 +1,7 @@ +from __future__ import annotations + import logging from datetime import datetime, timezone -from typing import Optional import pytest import pytz @@ -12,7 +13,7 @@ logger = logging.getLogger() class AwareDatetimeISOModel(BaseModel): - dt_optional: Optional[AwareDatetimeISO] = Field(default=None) + dt_optional: AwareDatetimeISO | None = Field(default=None) dt: AwareDatetimeISO diff --git a/tests/models/thl/test_product.py b/tests/models/thl/test_product.py index 5e9b249..39469dc 100644 --- a/tests/models/thl/test_product.py +++ b/tests/models/thl/test_product.py @@ -1,8 +1,10 @@ +from __future__ import annotations + import os import shutil from datetime import datetime, timedelta, timezone from decimal import Decimal -from typing import Callable, Optional +from typing import Callable from uuid import uuid4 import pytest @@ -10,7 +12,7 @@ from dask.distributed import Client as DaskClient from pydantic import ValidationError from generalresearch.currency import USDCent -from generalresearch.incite import GRLDatasets +from generalresearch.incite.base import GRLDatasets from generalresearch.incite.mergers.pop_ledger import PopLedgerMerge from generalresearch.managers.thl.ledger_manager.thl_ledger import ( ThlLedgerManager, @@ -25,6 +27,7 @@ from generalresearch.models.thl.product import ( IntegrationMode, PayoutConfig, PayoutTransformation, + PayoutTransformationPercentArgs, Product, ProfilingConfig, SourceConfig, @@ -137,6 +140,12 @@ class TestProduct: redirect_url="https://www.google.com/hey", ) + assert isinstance(p.payout_config.payout_transformation, PayoutTransformation) + assert isinstance( + p.payout_config.payout_transformation.kwargs, + PayoutTransformationPercentArgs, + ) + p.payout_config.payout_transformation = PayoutTransformation.model_validate( { "f": "payout_transformation_percent", @@ -576,7 +585,7 @@ class TestGlobalProductConfigFor: class TestProductFinancials: @pytest.fixture - def start(self) -> "datetime": + def start(self) -> datetime: return datetime(year=2018, month=3, day=14, hour=0, tzinfo=timezone.utc) @pytest.fixture @@ -584,7 +593,7 @@ class TestProductFinancials: return "30d" @pytest.fixture - def duration(self) -> Optional["timedelta"]: + def duration(self) -> timedelta | None: return None def test_balance( @@ -759,7 +768,7 @@ class TestProductFinancials: class TestProductBalance: @pytest.fixture - def start(self) -> "datetime": + def start(self) -> datetime: return datetime(year=2018, month=3, day=14, hour=0, tzinfo=timezone.utc) @pytest.fixture @@ -767,7 +776,7 @@ class TestProductBalance: return "30d" @pytest.fixture - def duration(self) -> Optional["timedelta"]: + def duration(self) -> timedelta | None: return None def test_inconsistent( @@ -783,7 +792,7 @@ class TestProductBalance: user_factory: Callable[..., User], session_with_tx_factory: Callable[..., Session], pop_ledger_merge, - start, + start: datetime, bp_payout_factory, payout_event_manager, ): @@ -792,8 +801,6 @@ class TestProductBalance: create_main_accounts() delete_df_collection(coll=ledger_collection) - from generalresearch.models.thl.user import User - u1: User = user_factory(product=product) # 1. Complete and Build Parquets 1st time @@ -827,7 +834,7 @@ class TestProductBalance: def test_not_inconsistent( self, product: Product, - mnt_filepath, + mnt_filepath: GRLDatasets, thl_lm: ThlLedgerManager, client_no_amm: DaskClient, delete_ledger_db, @@ -852,8 +859,6 @@ class TestProductBalance: create_main_accounts() delete_df_collection(coll=ledger_collection) - from generalresearch.models.thl.user import User - u1: User = user_factory(product=product) # 1. Complete and Build Parquets 1st time @@ -886,7 +891,7 @@ class TestProductBalance: class TestProductPOPFinancial: @pytest.fixture - def start(self) -> "datetime": + def start(self) -> datetime: return datetime(year=2018, month=3, day=14, hour=0, tzinfo=timezone.utc) @pytest.fixture @@ -894,15 +899,15 @@ class TestProductPOPFinancial: return "30d" @pytest.fixture - def duration(self) -> Optional["timedelta"]: + def duration(self) -> timedelta | None: return None def test_base( self, - product, - mnt_filepath, + product: Product, + mnt_filepath: GRLDatasets, thl_lm: ThlLedgerManager, - client_no_amm, + client_no_amm: DaskClient, delete_ledger_db, create_main_accounts, delete_df_collection, @@ -923,8 +928,6 @@ class TestProductPOPFinancial: create_main_accounts() delete_df_collection(coll=ledger_collection) - from generalresearch.models.thl.user import User - u1: User = user_factory(product=product) # 1. Complete and Build Parquets 1st time @@ -961,7 +964,7 @@ class TestProductPOPFinancial: class TestProductCache: @pytest.fixture - def start(self) -> "datetime": + def start(self) -> datetime: return datetime(year=2018, month=3, day=14, hour=0, tzinfo=timezone.utc) @pytest.fixture @@ -969,7 +972,7 @@ class TestProductCache: return "30d" @pytest.fixture - def duration(self) -> Optional["timedelta"]: + def duration(self) -> timedelta | None: return None def test_basic( @@ -1008,7 +1011,6 @@ class TestProductCache: ) from generalresearch.models.thl.product import Product - from generalresearch.models.thl.user import User u1: User = user_factory(product=product) @@ -1047,7 +1049,7 @@ class TestProductCache: def test_neg_balance_cache( self, product: Product, - mnt_filepath, + mnt_filepath: GRLDatasets, thl_lm, client_no_amm: DaskClient, thl_redis_config, @@ -1070,7 +1072,6 @@ class TestProductCache: delete_df_collection(coll=ledger_collection) from generalresearch.models.thl.product import Product - from generalresearch.models.thl.user import User u1: User = user_factory(product=product) @@ -1112,10 +1113,12 @@ class TestProductCache: # Fetch from cache and assert the instance loaded from redis rc = thl_redis_config.create_redis_client() - res: Optional[str] = rc.get(product.cache_key) + res: str | None = rc.get(product.cache_key) assert isinstance(res, str) p1: Product = Product.model_validate_json(res) + assert p1.balance + assert p1.balance.product_id == product.uuid assert p1.balance.payout_usd_str == "$0.71" assert p1.balance.adjustment == -71 |