Fix tests

Author: Val Brodsky

libs/labelbox/pyproject.toml

@@ -91,6 +91,7 @@ unit = "pytest tests/unit"
 integration = { cmd = "pytest tests/integration" }
 data = { cmd = "pytest tests/data" }
 rye-fmt-check = "rye fmt --check"
+MYPYPATH="../lbox-clients/src/"
 mypy-lint = "mypy src --pretty --show-error-codes --non-interactive --install-types"
 lint = { chain = ["mypy-lint", "rye-fmt-check"] }
 test = { chain = ["lint", "unit", "integration"] }

libs/labelbox/src/labelbox/client.py

@@ -202,9 +202,7 @@ def upload_file(self, path: str) -> str:
             )
 
     @retry.Retry(
-        predicate=retry.if_exception_type(
-            lbox.exceptions.InternalServerError
-        )
+        predicate=retry.if_exception_type(lbox.exceptions.InternalServerError)
     )
     def upload_data(
         self,
@@ -310,9 +308,7 @@ def _get_single(self, db_object_type, uid):
         res = self.execute(query_str, params)
         res = res and res.get(utils.camel_case(db_object_type.type_name()))
         if res is None:
-            raise lbox.exceptions.ResourceNotFoundError(
-                db_object_type, params
-            )
+            raise lbox.exceptions.ResourceNotFoundError(db_object_type, params)
         else:
             return db_object_type(self, res)
 
@@ -1975,9 +1971,7 @@ def is_feature_schema_archived(
                 )
 
         elif response.status_code == 404:
-            raise lbox.exceptions.ResourceNotFoundError(
-                Ontology, ontology_id
-            )
+            raise lbox.exceptions.ResourceNotFoundError(Ontology, ontology_id)
         else:
             raise lbox.exceptions.LabelboxError(
                 "Failed to get the feature schema archived status."
@@ -2006,9 +2000,7 @@ def get_model_slice(self, slice_id) -> ModelSlice:
         """
         res = self.execute(query_str, {"id": slice_id})
         if res is None or res["getSavedQuery"] is None:
-            raise lbox.exceptions.ResourceNotFoundError(
-                ModelSlice, slice_id
-            )
+            raise lbox.exceptions.ResourceNotFoundError(ModelSlice, slice_id)
 
         return Entity.ModelSlice(self, res["getSavedQuery"])
 
@@ -2316,9 +2308,7 @@ def get_embedding_by_name(self, name: str) -> Embedding:
         for e in embeddings:
             if e.name == name:
                 return e
-        raise lbox.exceptions.ResourceNotFoundError(
-            Embedding, dict(name=name)
-        )
+        raise lbox.exceptions.ResourceNotFoundError(Embedding, dict(name=name))
 
     def upsert_label_feedback(
         self, label_id: str, feedback: str, scores: Dict[str, float]
@@ -2365,8 +2355,7 @@ def upsert_label_feedback(
         scores_raw = res["upsertAutoQaLabelFeedback"]["scores"]
 
         return [
-            labelbox.LabelScore(name=x["name"], score=x["score"])
-            for x in scores_raw
+            LabelScore(name=x["name"], score=x["score"]) for x in scores_raw
         ]
 
     def get_labeling_service_dashboards(

libs/labelbox/src/labelbox/schema/annotation_import.py

@@ -15,9 +15,9 @@
     cast,
 )
 
-import lbox.exceptions
 import requests
 from google.api_core import retry
+from lbox.exceptions import ApiLimitError, NetworkError, ResourceNotFoundError
 from tqdm import tqdm  # type: ignore
 
 import labelbox
@@ -141,9 +141,9 @@ def wait_until_done(
 
     @retry.Retry(
         predicate=retry.if_exception_type(
-            lbox.exceptions.ApiLimitError,
-            lbox.exceptions.TimeoutError,
-            lbox.exceptions.NetworkError,
+            ApiLimitError,
+            TimeoutError,
+            NetworkError,
         )
     )
     def __backoff_refresh(self) -> None:
@@ -436,9 +436,7 @@ def from_name(
         }
         response = client.execute(query_str, params)
         if response is None:
-            raise lbox.exceptions.ResourceNotFoundError(
-                MEAPredictionImport, params
-            )
+            raise ResourceNotFoundError(MEAPredictionImport, params)
         response = response["modelErrorAnalysisPredictionImport"]
         if as_json:
             return response
@@ -561,9 +559,7 @@ def from_name(
         }
         response = client.execute(query_str, params)
         if response is None:
-            raise lbox.exceptions.ResourceNotFoundError(
-                MALPredictionImport, params
-            )
+            raise ResourceNotFoundError(MALPredictionImport, params)
         response = response["meaToMalPredictionImport"]
         if as_json:
             return response
@@ -710,9 +706,7 @@ def from_name(
         }
         response = client.execute(query_str, params)
         if response is None:
-            raise lbox.exceptions.ResourceNotFoundError(
-                MALPredictionImport, params
-            )
+            raise ResourceNotFoundError(MALPredictionImport, params)
         response = response["modelAssistedLabelingPredictionImport"]
         if as_json:
             return response
@@ -886,7 +880,7 @@ def from_name(
         }
         response = client.execute(query_str, params)
         if response is None:
-            raise lbox.exceptions.ResourceNotFoundError(LabelImport, params)
+            raise ResourceNotFoundError(LabelImport, params)
         response = response["labelImport"]
         if as_json:
             return response

libs/labelbox/src/labelbox/schema/user_group.py

@@ -7,7 +7,7 @@
     ResourceCreationError,
     ResourceNotFoundError,
     UnprocessableEntityError,
-)  # type: ignore
+)
 from pydantic import BaseModel, ConfigDict
 
 from labelbox import Client

libs/labelbox/tests/integration/test_client_errors.py

@@ -11,17 +11,17 @@
 
 
 def test_missing_api_key():
-    key = os.environ.get(labelbox.client._LABELBOX_API_KEY, None)
+    key = os.environ.get(lbox.request_client._LABELBOX_API_KEY, None)
     if key is not None:
-        del os.environ[labelbox.client._LABELBOX_API_KEY]
+        del os.environ[lbox.request_client._LABELBOX_API_KEY]
 
     with pytest.raises(lbox.exceptions.AuthenticationError) as excinfo:
         labelbox.client.Client()
 
     assert excinfo.value.message == "Labelbox API key not provided"
 
     if key is not None:
-        os.environ[labelbox.client._LABELBOX_API_KEY] = key
+        os.environ[lbox.request_client._LABELBOX_API_KEY] = key
 
 
 def test_bad_key(rand_gen):
@@ -73,7 +73,7 @@ def test_resource_not_found_error(client):
 
 def test_network_error(client):
     client = labelbox.client.Client(
-        api_key=client.api_key, endpoint="not_a_valid_URL"
+        api_key=client._request_client.api_key, endpoint="not_a_valid_URL"
     )
 
     with pytest.raises(lbox.exceptions.NetworkError) as excinfo:

libs/labelbox/tests/integration/test_foundry.py

@@ -1,7 +1,8 @@
-import labelbox as lb
 import pytest
-from labelbox.schema.foundry.app import App
+from lbox.exceptions import LabelboxError, ResourceNotFoundError
 
+import labelbox as lb
+from labelbox.schema.foundry.app import App
 from labelbox.schema.foundry.foundry_client import FoundryClient
 
 # Yolo object detection model id
@@ -97,7 +98,7 @@ def test_get_app(foundry_client, app):
 
 
 def test_get_app_with_invalid_id(foundry_client):
-    with pytest.raises(lb.exceptions.ResourceNotFoundError):
+    with pytest.raises(ResourceNotFoundError):
         foundry_client._get_app("invalid-id")
 
 
@@ -144,7 +145,7 @@ def test_run_foundry_app_returns_model_run_id(
 def test_run_foundry_with_invalid_data_row_id(foundry_client, app, random_str):
     invalid_datarow_id = "invalid-global-key"
     data_rows = lb.GlobalKeys([invalid_datarow_id])
-    with pytest.raises(lb.exceptions.LabelboxError) as exception:
+    with pytest.raises(LabelboxError) as exception:
         foundry_client.run_app(
             model_run_name=f"test-app-with-invalid-datarow-id-{random_str}",
             data_rows=data_rows,
@@ -156,7 +157,7 @@ def test_run_foundry_with_invalid_data_row_id(foundry_client, app, random_str):
 def test_run_foundry_with_invalid_global_key(foundry_client, app, random_str):
     invalid_global_key = "invalid-global-key"
     data_rows = lb.GlobalKeys([invalid_global_key])
-    with pytest.raises(lb.exceptions.LabelboxError) as exception:
+    with pytest.raises(LabelboxError) as exception:
         foundry_client.run_app(
             model_run_name=f"test-app-with-invalid-global-key-{random_str}",
             data_rows=data_rows,