diff --git a/packages/examples/cvat/exchange-oracle/README.md b/packages/examples/cvat/exchange-oracle/README.md index 639423301d..7d70eeac82 100644 --- a/packages/examples/cvat/exchange-oracle/README.md +++ b/packages/examples/cvat/exchange-oracle/README.md @@ -73,5 +73,6 @@ Available at `/docs` route To run tests ``` -docker compose -f docker-compose.test.yml up --build test --attach test --exit-code-from test +docker compose -f docker-compose.test.yml up --build test --attach test --exit-code-from test && \ + docker compose -f docker-compose.test.yml down ``` \ No newline at end of file diff --git a/packages/examples/cvat/exchange-oracle/src/services/cvat.py b/packages/examples/cvat/exchange-oracle/src/services/cvat.py index ca0e980c1e..5b2bbc96e0 100644 --- a/packages/examples/cvat/exchange-oracle/src/services/cvat.py +++ b/packages/examples/cvat/exchange-oracle/src/services/cvat.py @@ -794,7 +794,8 @@ def count_jobs_by_escrow_address( def get_free_job( session: Session, - cvat_projects: list[int], + escrow_address: str, + chain_id: int, *, user_wallet_address: str, for_update: bool | ForUpdateParams = False, @@ -805,7 +806,11 @@ def get_free_job( return ( _maybe_for_update(session.query(Job), enable=for_update) .where( - Job.cvat_project_id.in_(cvat_projects), + Job.project.has( + (Project.escrow_address == escrow_address) + & (Project.chain_id == chain_id) + & (Project.status == ProjectStatuses.annotation) + ), Job.status == JobStatuses.new, ~Job.assignments.any( ( @@ -984,22 +989,27 @@ def get_user_assignments_in_cvat_projects( ) -def count_active_user_assignments( +def has_active_user_assignments( session: Session, wallet_address: int, - cvat_projects: list[int], -) -> int: - return ( + escrow_address: str, + chain_id: int, +) -> bool: + return session.query( session.query(Assignment) .where( - Assignment.job.has(Job.cvat_project_id.in_(cvat_projects)), + Assignment.job.has( + Job.project.has( + (Project.escrow_address == escrow_address) & (Project.chain_id == chain_id) + ) + ), Assignment.user_wallet_address == wallet_address, Assignment.status == AssignmentStatuses.created.value, Assignment.completed_at == None, utcnow() < Assignment.expires_at, ) - .count() - ) + .exists() + ).scalar() # Image diff --git a/packages/examples/cvat/exchange-oracle/src/services/exchange.py b/packages/examples/cvat/exchange-oracle/src/services/exchange.py index c1fec576f0..23a7b8e1cc 100644 --- a/packages/examples/cvat/exchange-oracle/src/services/exchange.py +++ b/packages/examples/cvat/exchange-oracle/src/services/exchange.py @@ -18,7 +18,7 @@ def __str__(self) -> str: ) -def create_assignment(escrow_address: str, chain_id: Networks, wallet_address: str) -> str | None: # noqa: ARG001 (don't we want to use chain_id for filter?) +def create_assignment(escrow_address: str, chain_id: Networks, wallet_address: str) -> str | None: with SessionLocal.begin() as session: user = get_or_404( cvat_service.get_user_by_id(session, wallet_address, for_update=True), @@ -26,44 +26,37 @@ def create_assignment(escrow_address: str, chain_id: Networks, wallet_address: s "user", ) - # There can be several projects under one escrow, we need any - project = cvat_service.get_project_by_escrow_address( + if cvat_service.has_active_user_assignments( session, - escrow_address, - status_in=[ - ProjectStatuses.annotation - ], # avoid unnecessary locking on completed projects - for_update=True, - ) - - if not project: - # Retry without a lock to check if the project doesn't exist - get_or_404( - cvat_service.get_project_by_escrow_address( - session, escrow_address, status_in=[ProjectStatuses.annotation] - ), - escrow_address, - "job", - ) - return None - - has_active_assignments = ( - cvat_service.count_active_user_assignments( - session, wallet_address=wallet_address, cvat_projects=[project.cvat_id] - ) - > 0 - ) - if has_active_assignments: + wallet_address=wallet_address, + escrow_address=escrow_address, + chain_id=chain_id.value, + ): raise UserHasUnfinishedAssignmentError( "The user already has an unfinished assignment in this project" ) + # TODO: Try to put into 1 request. SQLAlchemy generates 2 queries with simple + # .options(selectinload(Job.project)) + project = get_or_404( + cvat_service.get_project_by_escrow_address( + session, escrow_address, status_in=[ProjectStatuses.annotation] + ), + escrow_address, + "job", + ) + unassigned_job = cvat_service.get_free_job( session, - cvat_projects=[project.cvat_id], + escrow_address=escrow_address, + chain_id=chain_id.value, user_wallet_address=wallet_address, for_update=True, + # lock the job to be able to make a rollback if CVAT requests fail + # can potentially be optimized to make less DB requests + # and rely only on assignment expiration ) + if not unassigned_job: return None @@ -72,7 +65,12 @@ def create_assignment(escrow_address: str, chain_id: Networks, wallet_address: s wallet_address=user.wallet_address, cvat_job_id=unassigned_job.cvat_id, expires_at=utcnow() - + timedelta(seconds=get_default_assignment_timeout(TaskTypes(project.job_type))), + + timedelta( + seconds=get_default_assignment_timeout( + TaskTypes(project.job_type) + # TODO: need to update this if we have multiple job types per escrow + ) + ), ) cvat_service.touch(session, Job, [unassigned_job.id]) diff --git a/packages/examples/cvat/exchange-oracle/tests/integration/services/test_exchange.py b/packages/examples/cvat/exchange-oracle/tests/integration/services/test_exchange.py index 25d1e6f81b..7d074d12e9 100644 --- a/packages/examples/cvat/exchange-oracle/tests/integration/services/test_exchange.py +++ b/packages/examples/cvat/exchange-oracle/tests/integration/services/test_exchange.py @@ -8,13 +8,14 @@ from fastapi import HTTPException from pydantic import ValidationError -from src.core.types import AssignmentStatuses, JobStatuses, Networks +from src.core.types import AssignmentStatuses, JobStatuses, Networks, TaskTypes from src.db import SessionLocal from src.endpoints.serializers import serialize_job from src.models.cvat import Assignment, User from src.schemas import exchange as service_api from src.services.exchange import create_assignment +from tests.utils.constants import ESCROW_ADDRESS, WALLET_ADDRESS1, WALLET_ADDRESS2 from tests.utils.db_helper import ( create_job, create_project, @@ -32,7 +33,7 @@ def tearDown(self): def test_serialize_job(self): cvat_id = 1 - escrow_address = "0x86e83d346041E8806e352681f3F14549C0d2BC67" + escrow_address = ESCROW_ADDRESS cvat_project = create_project(self.session, escrow_address, cvat_id) self.session.commit() @@ -64,7 +65,7 @@ def test_serialize_task_invalid_project(self): def test_serialize_task_invalid_manifest(self): cvat_id = 1 - escrow_address = "0x86e83d346041E8806e352681f3F14549C0d2BC67" + escrow_address = ESCROW_ADDRESS cvat_project = create_project(self.session, escrow_address, cvat_id) self.session.commit() @@ -75,10 +76,8 @@ def test_serialize_task_invalid_manifest(self): serialize_job(cvat_project) def test_create_assignment(self): - cvat_project_1, _, cvat_job_1 = create_project_task_and_job( - self.session, "0x86e83d346041E8806e352681f3F14549C0d2BC67", 1 - ) - user_address = "0x86e83d346041E8806e352681f3F14549C0d2BC69" + cvat_project_1, _, cvat_job_1 = create_project_task_and_job(self.session, ESCROW_ADDRESS, 1) + user_address = WALLET_ADDRESS1 user = User( wallet_address=user_address, cvat_email="test@hmt.ai", @@ -87,15 +86,9 @@ def test_create_assignment(self): self.session.add(user) self.session.commit() - with ( - open("tests/utils/manifest.json") as data, - patch("src.endpoints.serializers.get_escrow_manifest") as mock_get_manifest, - patch("src.services.exchange.cvat_api"), - ): - manifest = json.load(data) - mock_get_manifest.return_value = manifest + with patch("src.services.exchange.cvat_api"): assignment_id = create_assignment( - cvat_project_1.escrow_address, cvat_project_1.chain_id, user_address + cvat_project_1.escrow_address, Networks(cvat_project_1.chain_id), user_address ) assignment = self.session.query(Assignment).filter_by(id=assignment_id).first() @@ -105,15 +98,13 @@ def test_create_assignment(self): assert assignment.status == AssignmentStatuses.created def test_create_assignment_many_jobs_1_completed(self): - cvat_project, _, cvat_job_1 = create_project_task_and_job( - self.session, "0x86e83d346041E8806e352681f3F14549C0d2BC67", 1 - ) + cvat_project, _, cvat_job_1 = create_project_task_and_job(self.session, ESCROW_ADDRESS, 1) cvat_job_1.status = JobStatuses.completed.value cvat_task_2 = create_task(self.session, 2, cvat_project.cvat_id) cvat_job_2 = create_job(self.session, 2, cvat_task_2.cvat_id, cvat_project.cvat_id) - user_address = "0x86e83d346041E8806e352681f3F14549C0d2BC69" + user_address = WALLET_ADDRESS1 user = User( wallet_address=user_address, cvat_email="test@hmt.ai", @@ -135,15 +126,9 @@ def test_create_assignment_many_jobs_1_completed(self): self.session.commit() - with ( - open("tests/utils/manifest.json") as data, - patch("src.endpoints.serializers.get_escrow_manifest") as mock_get_manifest, - patch("src.services.exchange.cvat_api"), - ): - manifest = json.load(data) - mock_get_manifest.return_value = manifest + with patch("src.services.exchange.cvat_api"): assignment_id = create_assignment( - cvat_project.escrow_address, cvat_project.chain_id, user_address + cvat_project.escrow_address, Networks(cvat_project.chain_id), user_address ) assignment = self.session.query(Assignment).filter_by(id=assignment_id).first() @@ -153,20 +138,18 @@ def test_create_assignment_many_jobs_1_completed(self): assert assignment.status == AssignmentStatuses.created def test_create_assignment_invalid_user_address(self): - cvat_project_1, _, _ = create_project_task_and_job( - self.session, "0x86e83d346041E8806e352681f3F14549C0d2BC67", 1 - ) + cvat_project_1, _, _ = create_project_task_and_job(self.session, ESCROW_ADDRESS, 1) self.session.commit() - with pytest.raises(HTTPException): + with pytest.raises(HTTPException, match="Can't find user"): create_assignment( cvat_project_1.escrow_address, - cvat_project_1.chain_id, + Networks(cvat_project_1.chain_id), "invalid_address", ) def test_create_assignment_invalid_project(self): - user_address = "0x86e83d346041E8806e352681f3F14549C0d2BC69" + user_address = WALLET_ADDRESS1 user = User( wallet_address=user_address, cvat_email="test@hmt.ai", @@ -175,14 +158,12 @@ def test_create_assignment_invalid_project(self): self.session.add(user) self.session.commit() - with pytest.raises(HTTPException): - create_assignment("1", Networks.localhost.value, user_address) + with pytest.raises(HTTPException, match="Can't find job"): + create_assignment("1", Networks.localhost, user_address) def test_create_assignment_unfinished_assignment(self): - _, _, cvat_job = create_project_task_and_job( - self.session, "0x86e83d346041E8806e352681f3F14549C0d2BC67", 1 - ) - user_address = "0x86e83d346041E8806e352681f3F14549C0d2BC69" + _, _, cvat_job = create_project_task_and_job(self.session, ESCROW_ADDRESS, 1) + user_address = WALLET_ADDRESS1 user = User( wallet_address=user_address, cvat_email="test@hmt.ai", @@ -200,23 +181,52 @@ def test_create_assignment_unfinished_assignment(self): self.session.commit() with ( - open("tests/utils/manifest.json") as data, - patch("src.endpoints.serializers.get_escrow_manifest") as mock_get_manifest, patch("src.services.exchange.cvat_api"), + pytest.raises(Exception, match="unfinished assignment"), ): - manifest = json.load(data) - mock_get_manifest.return_value = manifest + create_assignment(ESCROW_ADDRESS, Networks.localhost, user_address) - with pytest.raises(HTTPException): - create_assignment("1", Networks.localhost.value, user_address) + def test_create_assignment_has_expired_assignment_and_available_jobs(self): + escrow_address = ESCROW_ADDRESS + project1, _, cvat_job1 = create_project_task_and_job(self.session, escrow_address, 1) + project2, _, cvat_job2 = create_project_task_and_job(self.session, escrow_address, 2) + project1.job_type = TaskTypes.image_skeletons_from_boxes + project2.job_type = TaskTypes.image_skeletons_from_boxes + self.session.add_all([project1, project2]) - def test_create_assignment_no_available_jobs_completed_assignment(self): - cvat_project, _, cvat_job_1 = create_project_task_and_job( - self.session, "0x86e83d346041E8806e352681f3F14549C0d2BC67", 1 + user_address = WALLET_ADDRESS1 + user = User( + wallet_address=user_address, + cvat_email="test@hmt.ai", + cvat_id=1, ) + self.session.add(user) + + old_assignment = Assignment( + id=str(uuid.uuid4()), + user_wallet_address=user_address, + cvat_job_id=cvat_job1.cvat_id, + created_at=datetime.now() - timedelta(hours=1), + expires_at=datetime.now() - timedelta(minutes=1), + status=AssignmentStatuses.expired.value, + ) + self.session.add(old_assignment) + + self.session.commit() + + with patch("src.services.exchange.cvat_api"): + new_assignment_id = create_assignment(escrow_address, Networks.localhost, user_address) + + new_assignment = self.session.query(Assignment).filter_by(id=new_assignment_id).first() + assert new_assignment.cvat_job_id == cvat_job2.cvat_id # job1 was attempted already + assert new_assignment.user_wallet_address == user_address + assert new_assignment.status == AssignmentStatuses.created + + def test_create_assignment_no_available_jobs_completed_assignment(self): + cvat_project, _, cvat_job_1 = create_project_task_and_job(self.session, ESCROW_ADDRESS, 1) cvat_job_1.status = JobStatuses.completed.value - user_address1 = "0x86e83d346041E8806e352681f3F14549C0d2BC69" + user_address1 = WALLET_ADDRESS1 user1 = User( wallet_address=user_address1, cvat_email="test1@hmt.ai", @@ -224,7 +234,7 @@ def test_create_assignment_no_available_jobs_completed_assignment(self): ) self.session.add(user1) - user_address2 = "0x86e83d346041E8806e352681f3F14549C0d2BC70" + user_address2 = WALLET_ADDRESS2 user2 = User( wallet_address=user_address2, cvat_email="test2@hmt.ai", @@ -246,25 +256,17 @@ def test_create_assignment_no_available_jobs_completed_assignment(self): self.session.commit() - with ( - open("tests/utils/manifest.json") as data, - patch("src.endpoints.serializers.get_escrow_manifest") as mock_get_manifest, - patch("src.services.exchange.cvat_api"), - ): - manifest = json.load(data) - mock_get_manifest.return_value = manifest + with patch("src.services.exchange.cvat_api"): assignment_id = create_assignment( - cvat_project.escrow_address, cvat_project.chain_id, user_address2 + cvat_project.escrow_address, Networks(cvat_project.chain_id), user_address2 ) assert assignment_id == None def test_create_assignment_no_available_jobs_active_foreign_assignment(self): - cvat_project, _, cvat_job_1 = create_project_task_and_job( - self.session, "0x86e83d346041E8806e352681f3F14549C0d2BC67", 1 - ) + cvat_project, _, cvat_job_1 = create_project_task_and_job(self.session, ESCROW_ADDRESS, 1) - user_address1 = "0x86e83d346041E8806e352681f3F14549C0d2BC69" + user_address1 = WALLET_ADDRESS1 user1 = User( wallet_address=user_address1, cvat_email="test1@hmt.ai", @@ -272,7 +274,7 @@ def test_create_assignment_no_available_jobs_active_foreign_assignment(self): ) self.session.add(user1) - user_address2 = "0x86e83d346041E8806e352681f3F14549C0d2BC70" + user_address2 = WALLET_ADDRESS2 user2 = User( wallet_address=user_address2, cvat_email="test2@hmt.ai", @@ -290,27 +292,19 @@ def test_create_assignment_no_available_jobs_active_foreign_assignment(self): self.session.commit() - with ( - open("tests/utils/manifest.json") as data, - patch("src.endpoints.serializers.get_escrow_manifest") as mock_get_manifest, - patch("src.services.exchange.cvat_api"), - ): - manifest = json.load(data) - mock_get_manifest.return_value = manifest + with patch("src.services.exchange.cvat_api"): assignment_id = create_assignment( - cvat_project.escrow_address, cvat_project.chain_id, user_address2 + cvat_project.escrow_address, Networks(cvat_project.chain_id), user_address2 ) assert assignment_id == None def test_create_assignment_wont_reassign_job_to_previous_user(self): - cvat_project_1, _, cvat_job_1 = create_project_task_and_job( - self.session, "0x86e83d346041E8806e352681f3F14549C0d2BC67", 1 - ) + cvat_project_1, _, cvat_job_1 = create_project_task_and_job(self.session, ESCROW_ADDRESS, 1) cvat_job_1.status = JobStatuses.new.value # validated and rejected return to 'new' user = User( - wallet_address="0x86e83d346041E8806e352681f3F14549C0d2BC69", + wallet_address=WALLET_ADDRESS1, cvat_email="test@hmt.ai", cvat_id=1, ) @@ -330,32 +324,26 @@ def test_create_assignment_wont_reassign_job_to_previous_user(self): self.session.commit() - with ( - open("tests/utils/manifest.json") as data, - patch("src.endpoints.serializers.get_escrow_manifest") as mock_get_manifest, - patch("src.services.exchange.cvat_api"), - ): - manifest = json.load(data) - mock_get_manifest.return_value = manifest + with patch("src.services.exchange.cvat_api"): assignment_id = create_assignment( - cvat_project_1.escrow_address, cvat_project_1.chain_id, user.wallet_address + cvat_project_1.escrow_address, + Networks(cvat_project_1.chain_id), + user.wallet_address, ) assert assignment_id is None def test_create_assignment_can_assign_job_to_new_user(self): - cvat_project_1, _, cvat_job_1 = create_project_task_and_job( - self.session, "0x86e83d346041E8806e352681f3F14549C0d2BC67", 1 - ) + cvat_project_1, _, cvat_job_1 = create_project_task_and_job(self.session, ESCROW_ADDRESS, 1) cvat_job_1.status = JobStatuses.new.value # validated and rejected return to 'new' previous_user = User( - wallet_address="0x86e83d346041E8806e352681f3F14549C0d2BC69", + wallet_address=WALLET_ADDRESS1, cvat_email="previous@hmt.ai", cvat_id=1, ) new_user = User( - wallet_address="0x69e83d346041E8806e352681f3F14549C0d2BC42", + wallet_address=WALLET_ADDRESS2, cvat_email="new@hmt.ai", cvat_id=2, ) @@ -376,15 +364,11 @@ def test_create_assignment_can_assign_job_to_new_user(self): self.session.commit() - with ( - open("tests/utils/manifest.json") as data, - patch("src.endpoints.serializers.get_escrow_manifest") as mock_get_manifest, - patch("src.services.exchange.cvat_api"), - ): - manifest = json.load(data) - mock_get_manifest.return_value = manifest + with patch("src.services.exchange.cvat_api"): assignment_id = create_assignment( - cvat_project_1.escrow_address, cvat_project_1.chain_id, new_user.wallet_address + cvat_project_1.escrow_address, + Networks(cvat_project_1.chain_id), + new_user.wallet_address, ) assignment = self.session.get(Assignment, assignment_id) diff --git a/packages/examples/cvat/exchange-oracle/tests/utils/constants.py b/packages/examples/cvat/exchange-oracle/tests/utils/constants.py index 2a5082c255..a1d535a814 100644 --- a/packages/examples/cvat/exchange-oracle/tests/utils/constants.py +++ b/packages/examples/cvat/exchange-oracle/tests/utils/constants.py @@ -20,6 +20,9 @@ TOKEN_ADDRESS = "0x976EA74026E726554dB657fA54763abd0C3a0aa9" FACTORY_ADDRESS = "0x14dC79964da2C08b23698B3D3cc7Ca32193d9955" +WALLET_ADDRESS1 = "0x86e83d346041E8806e352681f3F14549C0d2BC69" +WALLET_ADDRESS2 = "0x86e83d346041E8806e352681f3F14549C0d2BC70" + DEFAULT_MANIFEST_URL = "http://host.docker.internal:9000/manifests/manifest.json" DEFAULT_HASH = "test"