test: check imports

etang-ai21 · etang-ai21 · commit dcc878ed31fd · 2024-02-28T09:50:13.000+02:00
diff --git a/ai21/clients/sagemaker/resources/sagemaker_completion.py b/ai21/clients/sagemaker/resources/sagemaker_completion.py
@@ -1,4 +1,4 @@
-from __future__ import annotations
+# from __future__ import annotations
 
 from typing import List, Dict
 
diff --git a/tests/unittests/test_imports.py b/tests/unittests/test_imports.py
@@ -0,0 +1,20 @@
+from ai21 import __all__
+
+EXPECTED_ALL = [
+    "AI21EnvConfig",
+    "AI21Client",
+    "AI21APIError",
+    "APITimeoutError",
+    "AI21Error",
+    "MissingApiKeyError",
+    "ModelPackageDoesntExistError",
+    "TooManyRequestsError",
+    "AI21BedrockClient",
+    "AI21SageMakerClient",
+    "BedrockModelID",
+    "SageMaker",
+]
+
+
+def test_all_imports() -> None:
+    assert sorted(EXPECTED_ALL) == sorted(__all__)

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-from __future__ import annotations`
	`1`	`+# from __future__ import annotations`
`2`	`2`
`3`	`3`	`from typing import List, Dict`
`4`	`4`