test(cli): tests for rewriter.

Author: Zhe Yu

src/vectorcode/rewriter/__init__.py

@@ -11,7 +11,7 @@
 __all__ = ["RewriterBase", "OpenAIRewriter"]
 
 
-class RewriterError(Exception):
+class RewriterError(Exception):  # pragma: nocover
     pass
 
 
@@ -21,8 +21,9 @@ def get_rewriter(configs: Config) -> Optional[RewriterBase]:
         return None
     if configs.rewriter == "RewriterBase":
         raise RewriterError("RewriterBase is not a valid rewriter!")
-    rewriter_cls = getattr(sys.modules[__name__], configs.rewriter)
-    if issubclass(rewriter_cls, RewriterBase):
-        logger.info(f"Loaded {configs.rewriter}")
-        return rewriter_cls(configs)
+    if hasattr(sys.modules[__name__], configs.rewriter):
+        rewriter_cls = getattr(sys.modules[__name__], configs.rewriter)
+        if issubclass(rewriter_cls, RewriterBase):
+            logger.info(f"Loaded {configs.rewriter}")
+            return rewriter_cls(configs)
     raise RewriterError(f"Failed to find {configs.rewriter}!")

src/vectorcode/rewriter/base.py

@@ -3,7 +3,7 @@
 from vectorcode.cli_utils import Config
 
 
-class RewriterBase(ABC):
+class RewriterBase(ABC):  # pragma: nocover
     def __init__(self, config: Config) -> None:
         super().__init__()
         self.config = config

src/vectorcode/rewriter/openai.py

@@ -1,5 +1,4 @@
 import logging
-from typing import override
 
 import openai
 from openai.types.chat import ChatCompletion
@@ -49,7 +48,6 @@ def __init__(self, config: Config) -> None:
     No hallucinations.
 """
 
-    @override
     async def rewrite(self, original_query: list[str]):
         comp: ChatCompletion = self.client.beta.chat.completions.parse(
             messages=[

tests/rewriter/test_rewriter.py

@@ -0,0 +1,31 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from vectorcode.cli_utils import Config
+from vectorcode.rewriter import RewriterError, get_rewriter
+
+
+def test_get_rewriter():
+    assert get_rewriter(Config()) is None
+
+
+def test_get_rewriter_base():
+    with pytest.raises(RewriterError):
+        get_rewriter(Config(rewriter="RewriterBase"))
+
+
+def test_get_openai_rewriter():
+    with (
+        patch("vectorcode.rewriter.OpenAIRewriter") as mock_openai_cls,
+        patch("vectorcode.rewriter.issubclass") as mock_issubclass,
+    ):
+        mock_rewriter = MagicMock()
+        mock_openai_cls.return_value = mock_rewriter
+        mock_issubclass.return_value = True
+        assert get_rewriter(Config(rewriter="OpenAIRewriter")) == mock_rewriter
+
+
+def test_get_faulty_rewriter():
+    with pytest.raises(RewriterError):
+        get_rewriter(Config(rewriter="DummyRewriter"))