From 47e0f7797454e5058052352bf5627ff2f7fd038e Mon Sep 17 00:00:00 2001
From: etang <etang@ai21.com>
Date: Wed, 28 Feb 2024 09:50:13 +0200
Subject: [PATCH] test: check imports

---
 tests/unittests/test_imports.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)
 create mode 100644 tests/unittests/test_imports.py

diff --git a/tests/unittests/test_imports.py b/tests/unittests/test_imports.py
new file mode 100644
index 00000000..cb2ba4ed
--- /dev/null
+++ b/tests/unittests/test_imports.py
@@ -0,0 +1,23 @@
+# The following line should not be removed, as it is used to test the imports in runtime
+# noinspection PyUnresolvedReferences
+from ai21 import *  # noqa: F403
+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__)