from datetime import timezone, datetime import pytest from jb.models.event import MTurkEvent from generalresearchutils.pg_helper import PostgresConfig from datetime import datetime, timezone, timedelta from typing import Optional, TYPE_CHECKING, Callable, Generator from jb.managers.amt import AMTManager from jb.models.assignment import AssignmentStub, Assignment from generalresearchutils.currency import USDCent from jb.models.definitions import HitStatus, HitReviewStatus, AssignmentStatus from jb.models.hit import HitType, HitQuestion, Hit from tests import generate_amt_id if TYPE_CHECKING: from jb.managers.hit import HitQuestionManager, HitTypeManager, HitManager from jb.managers.assignment import AssignmentManager # --- MTurk Event --- @pytest.fixture def mturk_event( amt_assignment_id: str, amt_hit_id: str, amt_hit_type_id: str ) -> MTurkEvent: now = datetime.now(tz=timezone.utc) return MTurkEvent( event_type="AssignmentSubmitted", event_timestamp=now, amt_assignment_id=amt_assignment_id, amt_hit_type_id=amt_hit_type_id, amt_hit_id=amt_hit_id, ) # --- Question --- @pytest.fixture def question() -> HitQuestion: return HitQuestion(url="https://jamesbillings67.com/work/", height=1200) @pytest.fixture def question_record(hqm: "HitQuestionManager", question: HitQuestion) -> HitQuestion: return hqm.get_or_create(question) # --- HITType --- @pytest.fixture def hit_type() -> HitType: return HitType( title="Awesome Surveys!", description="Give us your opinion", reward=USDCent(5), keywords="market,research,amazing", min_active=10, ) @pytest.fixture def hit_type_record( pg_config: PostgresConfig, htm: "HitTypeManager", hit_type: HitType ) -> Generator[HitType, None, None]: hit_type.amt_hit_type_id = generate_amt_id() ht = htm.get_or_create(hit_type) yield ht with pg_config.make_connection() as conn: with conn.cursor() as c: c.execute("DELETE FROM mtwerk_hittype WHERE id=%s", (ht.id,)) conn.commit() @pytest.fixture def hit_type_record_with_amt_id( pg_config: PostgresConfig, htm: "HitTypeManager", hit_type: HitType ) -> Generator[HitType, None, None]: # This is a real hit type I've previously registered with amt (sandbox). # It will always exist hit_type.amt_hit_type_id = "3217B3DC4P5YW9DRV9R3X8O56V041J" # Get or create our db ht = htm.get_or_create(hit_type) yield ht with pg_config.make_connection() as conn: with conn.cursor() as c: c.execute("DELETE FROM mtwerk_hittype WHERE id=%s", (ht.id,)) conn.commit() # --- HIT --- @pytest.fixture def hit( amt_hit_id: str, amt_hit_type_id: str, amt_group_id: str, question: HitQuestion ) -> Hit: now = datetime.now(tz=timezone.utc) return Hit.model_validate( dict( amt_hit_id=amt_hit_id, amt_hit_type_id=amt_hit_type_id, amt_group_id=amt_group_id, status=HitStatus.Assignable, review_status=HitReviewStatus.NotReviewed, creation_time=now, expiration=now + timedelta(days=3), hit_question_xml=question.xml, qualification_requirements=[], max_assignments=1, assignment_pending_count=0, assignment_available_count=1, assignment_completed_count=0, description="Description", keywords="Keywords", reward=USDCent(5), title="Title", question_id=question.id, hit_type_id=None, ) ) @pytest.fixture def hit_record( pg_config: PostgresConfig, hm: "HitManager", question_record: HitQuestion, hit_type_record: HitType, hit: Hit, amt_hit_id: str, ) -> Generator[Hit, None, None]: """ Returns a hit that exists in our db, but does not in amazon (the amt ids are random). The mtwerk_hittype and mtwerk_question records will also exist (in the db) """ hit.hit_type_id = hit_type_record.id hit.amt_hit_id = amt_hit_id hit.question_id = question_record.id hm.create(hit) yield hit with pg_config.make_connection() as conn: with conn.cursor() as c: c.execute("DELETE FROM mtwerk_hit WHERE id=%s", (hit.id,)) conn.commit() @pytest.fixture def hit_in_amt( hm: "HitManager", question_record: HitQuestion, hit_type_record_with_amt_id: HitType ) -> Hit: # Actually create a new HIT in amt (sandbox) hit = AMTManager.create_hit_with_hit_type( hit_type=hit_type_record_with_amt_id, question=question_record ) # Create it in the DB hm.create(hit) return hit # --- Assignment --- @pytest.fixture def assignment_stub( hit: Hit, amt_assignment_id: str, amt_worker_id: str ) -> AssignmentStub: now = datetime.now(tz=timezone.utc) return AssignmentStub( amt_assignment_id=amt_assignment_id, amt_hit_id=hit.amt_hit_id, amt_worker_id=amt_worker_id, status=AssignmentStatus.Submitted, modified_at=now, created_at=now, ) @pytest.fixture def assignment_stub_record( pg_config: PostgresConfig, am: "AssignmentManager", hit_record: Hit, assignment_stub: AssignmentStub, ) -> Generator[AssignmentStub, None, None]: """ Returns an AssignmentStub that exists in our db, but does not in amazon (the amt ids are random). The mtwerk_hit, mtwerk_hittype, and mtwerk_question records will also exist (in the db) """ assignment_stub.hit_id = hit_record.id am.create_stub(stub=assignment_stub) yield assignment_stub with pg_config.make_connection() as conn: with conn.cursor() as c: c.execute( "DELETE FROM mtwerk_assignment WHERE id=%s", (assignment_stub.id,) ) conn.commit() @pytest.fixture def assignment(assignment_factory: Callable[..., Assignment]) -> Assignment: return assignment_factory() @pytest.fixture def assignment_record( pg_config: PostgresConfig, hit_record: Hit, assignment_record_factory: Callable[..., Assignment], ) -> Generator[Assignment, None, None]: assignment = assignment_record_factory(hit_id=hit_record.id) yield assignment with pg_config.make_connection() as conn: with conn.cursor() as c: c.execute( "DELETE FROM mtwerk_assignment WHERE id=%s", (assignment_stub.id,) ) conn.commit() @pytest.fixture def assignment_factory(hit_record: Hit) -> Callable[[Optional[str]], Assignment]: def _inner(amt_worker_id: Optional[str] = None) -> Assignment: now = datetime.now(tz=timezone.utc) amt_assignment_id = generate_amt_id() amt_worker_id = amt_worker_id or generate_amt_id() return Assignment( amt_assignment_id=amt_assignment_id, amt_hit_id=hit_record.amt_hit_id, amt_worker_id=amt_worker_id, status=AssignmentStatus.Submitted, modified_at=now, created_at=now, accept_time=now, auto_approval_time=now, submit_time=now, ) return _inner @pytest.fixture def assignment_record_factory( am: "AssignmentManager", assignment_factory: Callable[..., Assignment] ) -> Callable[..., Assignment]: def _inner(hit_id: int, amt_worker_id: Optional[str] = None) -> Assignment: a = assignment_factory(amt_worker_id=amt_worker_id) a.hit_id = hit_id am.create_stub(a) am.update_answer(a) return a return _inner