Less from init, more from generalresearch, basic db test from shared conftest

9 files changed, 56 insertions, 141 deletions

diff --git a/jb/models/__init__.py b/jb/models/__init__.py
index 7fe23a7..e69de29 100644
--- a/jb/models/__init__.py
+++ b/jb/models/__init__.py
@@ -1,39 +0,0 @@
-from decimal import Decimal
-
-from pydantic import BaseModel, ConfigDict, Field
-
-
-class HTTPHeaders(BaseModel):
-    request_id: str = Field(alias="x-amzn-requestid", min_length=36, max_length=36)
-    content_type: str = Field(alias="content-type", min_length=26, max_length=26)
-    # 'content-length': '1255',
-    content_length: str = Field(alias="content-length", min_length=2)
-    # 'Mon, 15 Jan 2024 23:40:32 GMT'
-    date: str = Field()
-
-    connection: str | None = Field(default=None)  # 'close'
-
-
-class ResponseMetadata(BaseModel):
-    model_config = ConfigDict(extra="forbid", validate_assignment=True)
-
-    request_id: str = Field(alias="RequestId", min_length=36, max_length=36)
-    status_code: int = Field(alias="HTTPStatusCode", ge=200, le=599)
-    headers: HTTPHeaders = Field(alias="HTTPHeaders")
-    retry_attempts: int = Field(alias="RetryAttempts", ge=0)
-
-
-class AMTAccount(BaseModel):
-    model_config = ConfigDict(extra="ignore", validate_assignment=True)
-
-    # Remaining available AWS Billing usage if you have enabled AWS Billing.
-    available_balance: Decimal = Field()
-    onhold_balance: Decimal = Field(default=Decimal(0))
-
-    # --- Properties ---
-
-    @property
-    def is_healthy(self) -> bool:
-        # A healthy account is one with at least $2,500 worth of
-        # credit available to it
-        return self.available_balance >= 2_500
diff --git a/jb/models/amt.py b/jb/models/amt.py
new file mode 100644
index 0000000..e012741
--- /dev/null
+++ b/jb/models/amt.py
@@ -0,0 +1,19 @@
+from decimal import Decimal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class AMTAccount(BaseModel):
+    model_config = ConfigDict(extra="ignore", validate_assignment=True)
+
+    # Remaining available AWS Billing usage if you have enabled AWS Billing.
+    available_balance: Decimal = Field()
+    onhold_balance: Decimal = Field(default=Decimal(0))
+
+    # --- Properties ---
+
+    @property
+    def is_healthy(self) -> bool:
+        # A healthy account is one with at least $2,500 worth of
+        # credit available to it
+        return self.available_balance >= 2_500
diff --git a/jb/models/assignment.py b/jb/models/assignment.py
index 1f7033d..fa6ccd5 100644
--- a/jb/models/assignment.py
+++ b/jb/models/assignment.py
@@ -1,7 +1,9 @@
+import logging
 from datetime import datetime, timezone
 from typing import Any, TypedDict
 from xml.etree import ElementTree
 
+from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr
 from mypy_boto3_mturk.type_defs import AssignmentTypeDef
 from pydantic import (
     BaseModel,
@@ -14,10 +16,11 @@ from pydantic import (
 )
 from typing_extensions import Self
 
-from jb.decorators import LOG
-from jb.models.custom_types import AMTBoto3ID, AwareDatetimeISO, UUIDStr
+from jb.models.custom_types import AMTBoto3ID
 from jb.models.definitions import AssignmentStatus
 
+logger = logging.getLogger("amtjb")
+
 
 class AnswerDict(TypedDict):
     amt_assignment_id: str
@@ -140,7 +143,7 @@ class Assignment(AssignmentStub):
                 values["tsid"] = TypeAdapter(UUIDStr).validate_python(tsid)
             except ValidationError as e:
                 # Don't break the model validation if a baddie messes with the tsid in the answer.
-                LOG.warning(e)
+                logger.warning(e)
 
                 values["tsid"] = None
         return values
diff --git a/jb/models/bonus.py b/jb/models/bonus.py
index 2c1d00c..c6da3c4 100644
--- a/jb/models/bonus.py
+++ b/jb/models/bonus.py
@@ -1,10 +1,11 @@
 from typing import Any
 
 from generalresearch.currency import USDCent
+from generalresearch.models.custom_types import AwareDatetimeISO, UUIDStr
 from pydantic import BaseModel, ConfigDict, Field, PositiveInt
 from typing_extensions import Self
 
-from jb.models.custom_types import AMTBoto3ID, AwareDatetimeISO, UUIDStr
+from jb.models.custom_types import AMTBoto3ID
 
 
 class Bonus(BaseModel):
diff --git a/jb/models/custom_types.py b/jb/models/custom_types.py
index a58dcb7..385c1ba 100644
--- a/jb/models/custom_types.py
+++ b/jb/models/custom_types.py
@@ -1,100 +1,8 @@
 import re
-from datetime import datetime, timezone
-from typing import Annotated, Any
-from uuid import UUID
+from typing import Annotated
 
-from pydantic import (
-    AwareDatetime,
-    HttpUrl,
-    StringConstraints,
-    TypeAdapter,
-)
-from pydantic.functional_serializers import PlainSerializer
-from pydantic.functional_validators import AfterValidator, BeforeValidator
-from pydantic.networks import UrlConstraints
-from pydantic_core import Url
-
-
-def convert_datetime_to_iso_8601_with_z_suffix(dt: datetime) -> str:
-    # By default, datetimes are serialized with the %f optional. We don't want that because
-    #   then the deserialization fails if the datetime didn't have microseconds.
-    return dt.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
-
-
-def convert_str_dt(v: Any) -> AwareDatetime | None:
-    # By default, pydantic is unable to handle tz-aware isoformat str. Attempt to parse a str
-    #   that was dumped using the iso8601 format with Z suffix.
-    if v is not None and type(v) is str:
-        assert v.endswith("Z") and "T" in v, "invalid format"
-        return datetime.strptime(v, "%Y-%m-%dT%H:%M:%S.%fZ").replace(
-            tzinfo=timezone.utc
-        )
-    return v
-
-
-def assert_utc(v: AwareDatetime) -> AwareDatetime:
-    assert v.tzinfo == timezone.utc, "Timezone is not UTC"
-    return v
-
-
-# Our custom AwareDatetime that correctly serializes and deserializes
-#   to an ISO8601 str with timezone
-AwareDatetimeISO = Annotated[
-    AwareDatetime,
-    BeforeValidator(convert_str_dt),
-    AfterValidator(assert_utc),
-    PlainSerializer(
-        lambda x: x.strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
-        when_used="json-unless-none",
-    ),
-]
-
-# ISO 3166-1 alpha-2 (two-letter codes, lowercase)
-# "Like" b/c it matches the format, but we're not explicitly checking
-#   it is one of our supported values. See models.thl.locales for that.
-CountryISOLike = Annotated[
-    str, StringConstraints(max_length=2, min_length=2, pattern=r"^[a-z]{2}$")
-]
-# 3-char ISO 639-2/B, lowercase
-LanguageISOLike = Annotated[
-    str, StringConstraints(max_length=3, min_length=3, pattern=r"^[a-z]{3}$")
-]
-
-
-def check_valid_uuid(v: str) -> str:
-    try:
-        assert UUID(v).hex == v
-    except Exception:
-        raise ValueError("Invalid UUID")
-    return v
-
-
-# Our custom field that stores a UUID4 as the .hex string representation
-UUIDStr = Annotated[
-    str,
-    StringConstraints(min_length=32, max_length=32),
-    AfterValidator(check_valid_uuid),
-]
-# Accepts the non-hex representation and coerces
-UUIDStrCoerce = Annotated[
-    str,
-    StringConstraints(min_length=32, max_length=32),
-    BeforeValidator(lambda value: TypeAdapter(UUID).validate_python(value).hex),
-    AfterValidator(check_valid_uuid),
-]
-
-# Same thing as UUIDStr with HttpUrl field. It is confusing that this
-# is not a str https://github.com/pydantic/pydantic/discussions/6395
-HttpUrlStr = Annotated[
-    str,
-    BeforeValidator(lambda value: str(TypeAdapter(HttpUrl).validate_python(value))),
-]
-
-HttpsUrl = Annotated[Url, UrlConstraints(max_length=2083, allowed_schemes=["https"])]
-HttpsUrlStr = Annotated[
-    str,
-    BeforeValidator(lambda value: str(TypeAdapter(HttpsUrl).validate_python(value))),
-]
+from pydantic import StringConstraints
+from pydantic.functional_validators import AfterValidator
 
 
 def check_valid_amt_boto3_id(v: str) -> str:
diff --git a/jb/models/errors.py b/jb/models/errors.py
index c590c6a..1fe71df 100644
--- a/jb/models/errors.py
+++ b/jb/models/errors.py
@@ -3,7 +3,7 @@ from enum import Enum
 
 from pydantic import BaseModel, ConfigDict, Field, model_validator
 
-from jb.models import ResponseMetadata
+from jb.models.response import ResponseMetadata
 
 
 class BotoRequestErrorOperation(str, Enum):
diff --git a/jb/models/event.py b/jb/models/event.py
index 0016ca7..fb5735b 100644
--- a/jb/models/event.py
+++ b/jb/models/event.py
@@ -1,9 +1,10 @@
 from typing import Any
 
+from generalresearch.models.custom_types import AwareDatetimeISO
 from mypy_boto3_mturk.literals import EventTypeType
 from pydantic import BaseModel, Field
 
-from jb.models.custom_types import AMTBoto3ID, AwareDatetimeISO
+from jb.models.custom_types import AMTBoto3ID
 
 
 class MTurkEvent(BaseModel):
diff --git a/jb/models/hit.py b/jb/models/hit.py
index a091c83..a550943 100644
--- a/jb/models/hit.py
+++ b/jb/models/hit.py
@@ -4,6 +4,7 @@ from uuid import uuid4
 from xml.etree import ElementTree
 
 from generalresearch.currency import USDCent
+from generalresearch.models.custom_types import AwareDatetimeISO, HttpsUrlStr
 from mypy_boto3_mturk.type_defs import HITTypeDef
 from pydantic import (
     BaseModel,
@@ -14,7 +15,7 @@ from pydantic import (
 )
 from typing_extensions import Self
 
-from jb.models.custom_types import AMTBoto3ID, AwareDatetimeISO, HttpsUrlStr
+from jb.models.custom_types import AMTBoto3ID
 from jb.models.definitions import HitReviewStatus, HitStatus
 
 
diff --git a/jb/models/response.py b/jb/models/response.py
new file mode 100644
index 0000000..22985af
--- /dev/null
+++ b/jb/models/response.py
@@ -0,0 +1,21 @@
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class HTTPHeaders(BaseModel):
+    request_id: str = Field(alias="x-amzn-requestid", min_length=36, max_length=36)
+    content_type: str = Field(alias="content-type", min_length=26, max_length=26)
+    # 'content-length': '1255',
+    content_length: str = Field(alias="content-length", min_length=2)
+    # 'Mon, 15 Jan 2024 23:40:32 GMT'
+    date: str = Field()
+
+    connection: str | None = Field(default=None)  # 'close'
+
+
+class ResponseMetadata(BaseModel):
+    model_config = ConfigDict(extra="forbid", validate_assignment=True)
+
+    request_id: str = Field(alias="RequestId", min_length=36, max_length=36)
+    status_code: int = Field(alias="HTTPStatusCode", ge=200, le=599)
+    headers: HTTPHeaders = Field(alias="HTTPHeaders")
+    retry_attempts: int = Field(alias="RetryAttempts", ge=0)