add a magic token flow specifically for amt account link. gr api manager add more logging and error handling

6 files changed, 69 insertions, 11 deletions

diff --git a/tests/conftest.py b/tests/conftest.py
index 25eb457..c056821 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -3,7 +3,7 @@ from typing import TYPE_CHECKING
 from uuid import uuid4
 from dotenv import load_dotenv
 import pytest
-from generalresearchutils.pg_helper import PostgresConfig
+from generalresearch.pg_helper import PostgresConfig
 from tests import generate_amt_id
 from _pytest.config import Config
 from jb.decorators import CLIENT_CONFIG
@@ -99,7 +99,7 @@ def settings(env_file_path: str) -> "Settings":
 
 @pytest.fixture(scope="session")
 def redis(settings: "Settings"):
-    from generalresearchutils.redis_helper import RedisConfig
+    from generalresearch.redis_helper import RedisConfig
 
     redis_config = RedisConfig(
         dsn=settings.redis,
diff --git a/tests/fixtures/flow.py b/tests/fixtures/flow.py
index dd2f83e..08ec49e 100644
--- a/tests/fixtures/flow.py
+++ b/tests/fixtures/flow.py
@@ -4,9 +4,9 @@ from uuid import uuid4
 
 import pytest
 import requests
-from generalresearchutils.models.thl.payout import UserPayoutEvent
-from generalresearchutils.models.thl.wallet import PayoutType
-from generalresearchutils.models.thl.wallet.cashout_method import (
+from generalresearch.models.thl.payout import UserPayoutEvent
+from generalresearch.models.thl.wallet import PayoutType
+from generalresearch.models.thl.wallet.cashout_method import (
     CashoutRequestResponse,
     CashoutRequestInfo,
 )
@@ -20,8 +20,8 @@ from jb.managers.amt import (
     APPROVAL_MESSAGE,
     BONUS_MESSAGE,
 )
-from generalresearchutils.currency import USDCent
-from generalresearchutils.models.thl.definitions import PayoutStatus
+from generalresearch.currency import USDCent
+from generalresearch.models.thl.definitions import PayoutStatus
 
 
 @pytest.fixture
diff --git a/tests/fixtures/managers.py b/tests/fixtures/managers.py
index a3187d7..d10b542 100644
--- a/tests/fixtures/managers.py
+++ b/tests/fixtures/managers.py
@@ -1,7 +1,7 @@
 from typing import TYPE_CHECKING
 import pytest
 from jb.managers import Permission
-from generalresearchutils.pg_helper import PostgresConfig
+from generalresearch.pg_helper import PostgresConfig
 from mypy_boto3_mturk import MTurkClient
 
 if TYPE_CHECKING:
diff --git a/tests/fixtures/models.py b/tests/fixtures/models.py
index 671c7b3..b818caa 100644
--- a/tests/fixtures/models.py
+++ b/tests/fixtures/models.py
@@ -3,13 +3,13 @@ from datetime import timezone, datetime
 import pytest
 
 from jb.models.event import MTurkEvent
-from generalresearchutils.pg_helper import PostgresConfig
+from generalresearch.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 generalresearch.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
diff --git a/tests/flow/test_tasks.py b/tests/flow/test_tasks.py
index 3a71504..f939d20 100644
--- a/tests/flow/test_tasks.py
+++ b/tests/flow/test_tasks.py
@@ -16,7 +16,7 @@ from jb.managers.amt import (
 from mypy_boto3_mturk.type_defs import (
     GetAssignmentResponseTypeDef,
 )
-from generalresearchutils.currency import USDCent
+from generalresearch.currency import USDCent
 from jb.managers.assignment import AssignmentManager
 from jb.managers.bonus import BonusManager
 from jb.managers.hit import HitManager
diff --git a/tests/http/test_auth.py b/tests/http/test_auth.py
index 8a625bf..02ac88a 100644
--- a/tests/http/test_auth.py
+++ b/tests/http/test_auth.py
@@ -13,6 +13,8 @@ from jb.models.auth import User
 class FakeGRApiManager:
     def __init__(self):
         self.users: dict[str, User] = {}
+        self.amt_worker_ids: set[str] = set()
+        self.transition_calls: list[tuple[User, str]] = []
 
     def ensure_user_exists(self, user: User) -> User:
         self.users.setdefault(user.product_user_id, user)
@@ -21,6 +23,18 @@ class FakeGRApiManager:
     def get_user(self, product_user_id: str) -> User:
         return self.users[product_user_id]
 
+    def add_amt_user(self, amt_worker_id: str) -> None:
+        self.amt_worker_ids.add(amt_worker_id)
+
+    def transition_user_from_amt(self, user: User, amt_worker_id: str) -> User:
+        self.transition_calls.append((user, amt_worker_id))
+        if amt_worker_id not in self.amt_worker_ids:
+            raise ValueError(f"User {amt_worker_id} does not exist")
+
+        self.amt_worker_ids.remove(amt_worker_id)
+        self.users[user.product_user_id] = user
+        return user
+
 
 @pytest.fixture
 def email() -> str:
@@ -64,3 +78,47 @@ class TestAuth:
         res = await client.get("/auth/session")
         assert res.status_code == 200
         assert res.json()["email"] == email
+
+    @pytest.mark.anyio
+    async def test_amt_account_link(
+        self,
+        httpxclient: AsyncClient,
+        fake_gr_api_manager: FakeGRApiManager,
+        email: str,
+        amt_worker_id: str,
+    ):
+        client = httpxclient
+        fake_gr_api_manager.add_amt_user(amt_worker_id)
+
+        res = await client.post(
+            "/auth/link-amt/request",
+            json={"email": email, "amt_worker_id": amt_worker_id},
+        )
+        d = res.json()
+        assert res.status_code == 200
+        assert d["magic_link"]
+        assert fake_gr_api_manager.transition_calls == []
+
+        token = parse_qs(urlparse(d["magic_link"]).query)["token"][0]
+
+        url = "/auth/link-amt/exchange"
+        body = {"token": token}
+        res = await client.post(url, json=body)
+        assert res.status_code == 204
+        assert client.cookies.get(SESSION_COOKIE_NAME)
+        assert len(fake_gr_api_manager.transition_calls) == 1
+        transitioned_user, transitioned_amt_worker_id = (
+            fake_gr_api_manager.transition_calls[0]
+        )
+        assert transitioned_user.email == email
+        assert transitioned_amt_worker_id == amt_worker_id
+        assert amt_worker_id not in fake_gr_api_manager.amt_worker_ids
+
+        res = await client.get("/auth/session")
+        assert res.status_code == 200
+        assert res.json()["email"] == email
+
+        # Make sure we can't do it again
+        res = await client.post(url, json=body)
+        assert res.status_code == 401
+        assert len(fake_gr_api_manager.transition_calls) == 1