Fix the unit tests

tests/fixtures.py

@@ -1,14 +1,13 @@
 import uuid
 from datetime import datetime
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
 
 from encord import EncordUserClient, Project
 from encord.client import EncordClientProject
-from encord.http.querier import Querier
 from encord.ontology import Ontology
 from encord.orm.ontology import Ontology as OrmOntology
 from encord.orm.project import ProjectDTO, ProjectType
@@ -25,7 +24,7 @@
 
 @pytest.fixture
 def ontology() -> Ontology:
-    return Ontology(None, OrmOntology.from_dict(ONTOLOGY_BLURB), None)
+    return Ontology(OrmOntology.from_dict(ONTOLOGY_BLURB), MagicMock())
 
 
 @pytest.fixture
@@ -35,9 +34,14 @@ def user_client() -> EncordUserClient:
 
 @pytest.fixture
 @patch.object(EncordClientProject, "get_project_v2")
-@patch.object(Querier, "basic_getter")
-def project(querier_mock: Querier, client_project_mock, user_client: EncordUserClient, ontology: Ontology) -> Project:
-    querier_mock.return_value = OrmOntology.from_dict(ONTOLOGY_BLURB)
+@patch.object(EncordUserClient, "get_ontology")
+def project(
+    client_ontology_mock: MagicMock,
+    client_project_mock: MagicMock,
+    user_client: EncordUserClient,
+    ontology: Ontology,
+) -> Project:
+    client_ontology_mock.return_value = ontology
 
     client_project_mock.return_value = ProjectDTO(
         project_hash=uuid.uuid4(),

tests/http/test_timeout_overrides_setting.py

@@ -40,7 +40,18 @@ def stub_responses(*args, **kwargs) -> MagicMock:
         }
         mock_response.status_code = 200
         mock_response.json.return_value = response
-
+    elif args[0].path_url.startswith("/v2/public/ontologies"):
+        response = {
+            "ontologyUuid": ONTOLOGY_HASH,
+            "title": "Test ontology",
+            "description": "",
+            "created_at": "2024-01-01 11:11:11",
+            "last_edited_at": "2024-01-01 11:11:11",
+            "editor": {"objects": [], "classifications": []},
+            "user_role": 1,
+        }
+        mock_response.status_code = 200
+        mock_response.json.return_value = response
     else:
         # Legacy api
         body_json = json.loads(args[0].body)

tests/test_user_client_auth.py

@@ -15,45 +15,47 @@
 from encord.configs import SshConfig
 from encord.http.v2.payloads import Page
 from encord.orm.analytics import CollaboratorTimer
-from encord.orm.ontology import Ontology as OrmOntology
-from encord.orm.ontology import OntologyStructure
 from encord.orm.project import Project as OrmProject
 from encord.orm.project import ProjectDTO, ProjectType
 from encord.user_client import EncordUserClient
+from encord.utilities.ontology_user import OntologyUserRole, OntologyWithUserRole
 from tests.fixtures import PRIVATE_KEY_PEM
 
 PROJECT_HASH = str(uuid.uuid4())
-ONTOLOGY_HASH = str(uuid.uuid4())
+DATASET_HASH = str(uuid.uuid4())
+ONTOLOGY_UUID = uuid.uuid4()
 
 
-ontology_orm = OrmOntology(
-    ontology_hash=ONTOLOGY_HASH,
-    title="Test ontology",
-    structure=OntologyStructure(),
-    created_at=datetime.now(),
-    last_edited_at=datetime.now(),
-)
-
-ontology_dic = {
-    "ontology_hash": ONTOLOGY_HASH,
-    "title": "Test ontology",
+dataset_dic = {
+    "dataset_hash": DATASET_HASH,
+    "dataset_type": 0,
+    "title": "Test dataset",
     "created_at": datetime.now(),
     "last_edited_at": datetime.now(),
-    "editor": OntologyStructure().to_dict(),
     "description": "",
 }
 
+ontology_orm = OntologyWithUserRole(
+    ontology_uuid=ONTOLOGY_UUID,
+    title="Dummy ontology",
+    description="",
+    editor={"objects": [], "classifications": []},
+    created_at=datetime.now(),
+    last_edited_at=datetime.now(),
+    user_role=OntologyUserRole.USER,
+)
+
 api_key_meta_dic = {"title": "dummy key", "resource_type": "project"}
 
-project_orm = OrmProject({"project_hash": PROJECT_HASH, "ontology_hash": ONTOLOGY_HASH})
+project_orm = OrmProject({"project_hash": PROJECT_HASH, "ontology_hash": DATASET_HASH})
 project_dto = ProjectDTO(
     project_hash=uuid.uuid4(),
     project_type=ProjectType.MANUAL_QA,
     title="Dummy project",
     description="",
     created_at=datetime.now(),
     last_edited_at=datetime.now(),
-    ontology_hash="dummy-ontology-hash",
+    ontology_hash=str(ONTOLOGY_UUID),
     workflow=None,
 )
 
@@ -63,7 +65,7 @@ def bearer_token() -> str:
     return "a-fake-token"
 
 
-def make_side_effects(project_response: Optional[MagicMock] = None, ontology_response: Optional[MagicMock] = None):
+def make_side_effects(project_response: Optional[MagicMock] = None):
     # Mocking a few service endpoints to make it possible to test authentication
 
     def side_effects(*args, **kwargs):
@@ -87,13 +89,11 @@ def side_effects(*args, **kwargs):
                 mock_response.json.return_value = {"status": 200, "response": project_orm}
                 mock_response.content = json.dumps({"status": 200, "response": project_orm}, default=str)
                 return mock_response
-            if request_type == "ontology":
-                if ontology_response:
-                    return ontology_response
+            if request_type == "dataset":
                 mock_response = MagicMock()
                 mock_response.status_code = 200
-                mock_response.json.return_value = {"status": 200, "response": ontology_dic}
-                mock_response.content = json.dumps({"status": 200, "response": ontology_dic}, default=str)
+                mock_response.json.return_value = {"status": 200, "response": dataset_dic}
+                mock_response.content = json.dumps({"status": 200, "response": dataset_dic}, default=str)
                 return mock_response
             else:
                 print(f"Unknown type {request_type}")
@@ -111,6 +111,13 @@ def side_effects(*args, **kwargs):
             mock_response.json.return_value = Page[CollaboratorTimer](results=[]).to_dict()
             mock_response.json.content = json.dumps(Page[CollaboratorTimer](results=[]).to_dict())
             return mock_response
+        elif args[0].method == "GET" and args[0].path_url.startswith("/v2/public/ontologies/"):
+            mock_response = MagicMock()
+            mock_response.status_code = 200
+
+            mock_response.json.return_value = ontology_orm.to_dict()
+            mock_response.content = json.dumps(ontology_orm.to_dict())
+            return mock_response
         else:
             raise RuntimeError(f"Not mocked URL: {args[0].path_url}")
 
@@ -131,16 +138,16 @@ def test_v1_public_resource_when_initialised_with_ssh_key(mock_send, bearer_toke
     mock_send.side_effect = make_side_effects()
 
     user_client = EncordUserClient.create_with_ssh_private_key(ssh_private_key=PRIVATE_KEY_PEM)
-    user_client.get_ontology(ONTOLOGY_HASH)
+    user_client.get_dataset(DATASET_HASH)
 
     assert mock_send.call_count == 1
     for mock_call in mock_send.call_args_list:
         # Expect call to have correct resource type and id, and correct bearer auth
         request = mock_call.args[0]
 
         assert request.path_url == "/public"
-        assert request.headers["ResourceType"] == "ontology"
-        assert request.headers["ResourceID"] == ONTOLOGY_HASH
+        assert request.headers["ResourceType"] == "dataset"
+        assert request.headers["ResourceID"] == DATASET_HASH
         assert request.headers["Authorization"] == get_encord_auth_header(request)
 
 
@@ -149,14 +156,14 @@ def test_v1_public_resource_when_initialised_with_bearer_auth(mock_send, bearer_
     mock_send.side_effect = make_side_effects()
 
     user_client = EncordUserClient.create_with_bearer_token(bearer_token)
-    user_client.get_ontology(ONTOLOGY_HASH)
+    user_client.get_dataset(DATASET_HASH)
 
     assert mock_send.call_count == 1
     for mock_call in mock_send.call_args_list:
         # Expect call to have correct resource type and id, and correct bearer auth
         assert mock_call.args[0].path_url == "/public"
-        assert mock_call.args[0].headers["ResourceType"] == "ontology"
-        assert mock_call.args[0].headers["ResourceID"] == ONTOLOGY_HASH
+        assert mock_call.args[0].headers["ResourceType"] == "dataset"
+        assert mock_call.args[0].headers["ResourceID"] == DATASET_HASH
         assert mock_call.args[0].headers["Authorization"] == f"Bearer {bearer_token}"