Support optional imports in unit tests (#293)

pinecone-io · Feb 14, 2024 · 181762c · 181762c
1 parent c9ce812
commit 181762c
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/tests/unit/tokenizer/test_cohere_hf_tokenizer.py b/tests/unit/tokenizer/test_cohere_hf_tokenizer.py
@@ -8,7 +8,14 @@ class TestCohereHFTokenizer(BaseTestTokenizer):
     @staticmethod
     @pytest.fixture(scope="class")
     def tokenizer():
-        return CohereHFTokenizer()
+        try:
+            tokenizer = CohereHFTokenizer()
+        except ImportError:
+            pytest.skip(
+                "`cohere` extra not installed. Skipping CohereHFTokenizer unit "
+                "tests"
+            )
+        return tokenizer
 
     @staticmethod
     @pytest.fixture

diff --git a/tests/unit/tokenizer/test_llama_tokenizer.py b/tests/unit/tokenizer/test_llama_tokenizer.py
@@ -8,7 +8,14 @@ class TestLlamaTokenizer(BaseTestTokenizer):
     @staticmethod
     @pytest.fixture(scope="class")
     def tokenizer():
-        return LlamaTokenizer(model_name="hf-internal-testing/llama-tokenizer")
+        try:
+            tokenizer = LlamaTokenizer(model_name="hf-internal-testing/llama-tokenizer")
+        except ImportError:
+            pytest.skip(
+                "`transformers` extra not installed. Skipping LLamaTokenizer unit "
+                "tests"
+            )
+        return tokenizer
 
     @staticmethod
     @pytest.fixture