Refactor ParamMixin to use StudyKeyStrategy

imednet/core/endpoint/mixins/params.py

@@ -2,7 +2,13 @@
 
 from typing import Any, Dict, Optional, cast
 
-from imednet.core.endpoint.strategies import DefaultParamProcessor
+from imednet.core.endpoint.strategies import (
+    DefaultParamProcessor,
+    KeepStudyKeyStrategy,
+    OptionalStudyKeyStrategy,
+    PopStudyKeyStrategy,
+    StudyKeyStrategy,
+)
 from imednet.core.endpoint.structs import ParamState
 from imednet.core.protocols import ParamProcessor
 from imednet.utils.filters import build_filter_string
@@ -18,6 +24,25 @@ class ParamMixin:
     _missing_study_exception: type[Exception] = ValueError
 
     PARAM_PROCESSOR_CLS: type[ParamProcessor] = DefaultParamProcessor
+    STUDY_KEY_STRATEGY: Optional[StudyKeyStrategy] = None
+
+    @property
+    def study_key_strategy(self) -> StudyKeyStrategy:
+        """
+        Get the configured study key strategy.
+
+        Returns:
+            The strategy instance to use.
+        """
+        if self.STUDY_KEY_STRATEGY:
+            return self.STUDY_KEY_STRATEGY
+
+        # Backward compatibility logic
+        if self.requires_study_key:
+            if self._pop_study_filter:
+                return PopStudyKeyStrategy(exception_cls=self._missing_study_exception)
+            return KeepStudyKeyStrategy(exception_cls=self._missing_study_exception)
+        return OptionalStudyKeyStrategy()
 
     def _resolve_params(
         self,
@@ -41,21 +66,8 @@ def _resolve_params(
         if study_key:
             filters["studyKey"] = study_key
 
-        study: Optional[str] = None
-        if self.requires_study_key:
-            if self._pop_study_filter:
-                try:
-                    study = filters.pop("studyKey")
-                except KeyError as exc:
-                    raise self._missing_study_exception(
-                        "Study key must be provided or set in the context"
-                    ) from exc
-            else:
-                study = filters.get("studyKey")
-                if not study:
-                    raise ValueError("Study key must be provided or set in the context")
-        else:
-            study = filters.get("studyKey")
+        # Delegate study key handling to the strategy
+        study, filters = self.study_key_strategy.process(filters)
 
         other_filters = {k: v for k, v in filters.items() if k != "studyKey"}
 

imednet/core/endpoint/strategies.py

@@ -5,7 +5,7 @@
 to customize how filters are processed and special parameters are extracted.
 """
 
-from typing import Any, Dict, Tuple
+from typing import Any, Dict, Optional, Protocol, Tuple, Type, runtime_checkable
 
 from imednet.core.protocols import ParamProcessor
 
@@ -28,3 +28,107 @@ def process_filters(self, filters: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict
             A tuple of (copy of filters, empty dict).
         """
         return filters.copy(), {}
+
+
+@runtime_checkable
+class StudyKeyStrategy(Protocol):
+    """Protocol for study key handling strategies."""
+
+    def process(self, filters: Dict[str, Any]) -> Tuple[Optional[str], Dict[str, Any]]:
+        """
+        Process the study key from filters.
+
+        Args:
+            filters: The current filters dictionary.
+
+        Returns:
+            A tuple containing the extracted study key (if any) and the modified filters.
+        """
+        ...
+
+
+class KeepStudyKeyStrategy:
+    """
+    Strategy that requires a study key and keeps it in the filters.
+
+    Used when the API expects 'studyKey' as a query parameter.
+    """
+
+    def __init__(self, exception_cls: Type[Exception] = ValueError) -> None:
+        self._exception_cls = exception_cls
+
+    def process(self, filters: Dict[str, Any]) -> Tuple[Optional[str], Dict[str, Any]]:
+        """
+        Extract study key, validate presence, and keep in filters.
+
+        Args:
+            filters: The filters dictionary.
+
+        Returns:
+            Tuple of (study_key, filters).
+
+        Raises:
+            Exception: If study key is missing (type determined by exception_cls).
+        """
+        study_key = filters.get("studyKey")
+        if not study_key:
+            raise self._exception_cls("Study key must be provided or set in the context")
+        return study_key, filters
+
+
+class PopStudyKeyStrategy:
+    """
+    Strategy that requires a study key but removes it from the filters.
+
+    Used when the study key is part of the path or handled separately,
+    not sent as a query parameter.
+    """
+
+    def __init__(self, exception_cls: Type[Exception] = ValueError) -> None:
+        self._exception_cls = exception_cls
+
+    def process(self, filters: Dict[str, Any]) -> Tuple[Optional[str], Dict[str, Any]]:
+        """
+        Extract study key, validate presence, and remove from filters.
+
+        Args:
+            filters: The filters dictionary.
+
+        Returns:
+            Tuple of (study_key, modified_filters).
+
+        Raises:
+            Exception: If study key is missing (type determined by exception_cls).
+        """
+        # Ensure we work on a copy if we are modifying it,
+        # but the mixin usually passes a copy or we should copy here.
+        # ParamProcessor returns a copy, so filters here might be that copy.
+        # But to be safe and pure:
+        filters_copy = filters.copy()
+
+        if "studyKey" not in filters_copy:
+            raise self._exception_cls("Study key must be provided or set in the context")
+
+        study_key = filters_copy.pop("studyKey")
+        return study_key, filters_copy
+
+
+class OptionalStudyKeyStrategy:
+    """
+    Strategy that allows the study key to be optional.
+
+    If present, it is kept in the filters.
+    """
+
+    def process(self, filters: Dict[str, Any]) -> Tuple[Optional[str], Dict[str, Any]]:
+        """
+        Extract study key if present.
+
+        Args:
+            filters: The filters dictionary.
+
+        Returns:
+            Tuple of (study_key_or_none, filters).
+        """
+        study_key = filters.get("studyKey")
+        return study_key, filters

tests/unit/test_study_key_strategies.py

@@ -0,0 +1,98 @@
+import pytest
+
+from imednet.core.endpoint.strategies import (
+    KeepStudyKeyStrategy,
+    OptionalStudyKeyStrategy,
+    PopStudyKeyStrategy,
+)
+
+
+class TestKeepStudyKeyStrategy:
+    def test_process_with_valid_key(self):
+        strategy = KeepStudyKeyStrategy()
+        filters = {"studyKey": "valid-key", "other": "param"}
+        key, new_filters = strategy.process(filters)
+
+        assert key == "valid-key"
+        assert "studyKey" in new_filters
+        assert new_filters["studyKey"] == "valid-key"
+        assert new_filters["other"] == "param"
+
+    def test_process_missing_key_raises_error(self):
+        strategy = KeepStudyKeyStrategy()
+        filters = {"other": "param"}
+
+        with pytest.raises(ValueError, match="Study key must be provided"):
+            strategy.process(filters)
+
+    def test_process_custom_exception(self):
+        class CustomError(Exception):
+            pass
+
+        strategy = KeepStudyKeyStrategy(exception_cls=CustomError)
+        filters = {"other": "param"}
+
+        with pytest.raises(CustomError):
+            strategy.process(filters)
+
+    def test_process_empty_key_raises_error(self):
+        strategy = KeepStudyKeyStrategy()
+        filters = {"studyKey": "", "other": "param"}
+
+        with pytest.raises(ValueError, match="Study key must be provided"):
+            strategy.process(filters)
+
+
+class TestPopStudyKeyStrategy:
+    def test_process_with_valid_key(self):
+        strategy = PopStudyKeyStrategy()
+        filters = {"studyKey": "valid-key", "other": "param"}
+        key, new_filters = strategy.process(filters)
+
+        assert key == "valid-key"
+        assert "studyKey" not in new_filters
+        assert new_filters["other"] == "param"
+
+    def test_process_missing_key_raises_error(self):
+        strategy = PopStudyKeyStrategy()
+        filters = {"other": "param"}
+
+        with pytest.raises(ValueError, match="Study key must be provided"):
+            strategy.process(filters)
+
+    def test_process_custom_exception(self):
+        class CustomError(Exception):
+            pass
+
+        strategy = PopStudyKeyStrategy(exception_cls=CustomError)
+        filters = {"other": "param"}
+
+        with pytest.raises(CustomError):
+            strategy.process(filters)
+
+    def test_process_empty_key_returns_empty_string(self):
+        strategy = PopStudyKeyStrategy()
+        filters = {"studyKey": "", "other": "param"}
+        key, new_filters = strategy.process(filters)
+        assert key == ""
+        assert "studyKey" not in new_filters
+
+
+class TestOptionalStudyKeyStrategy:
+    def test_process_with_key(self):
+        strategy = OptionalStudyKeyStrategy()
+        filters = {"studyKey": "valid-key", "other": "param"}
+        key, new_filters = strategy.process(filters)
+
+        assert key == "valid-key"
+        assert "studyKey" in new_filters
+        assert new_filters["studyKey"] == "valid-key"
+
+    def test_process_without_key(self):
+        strategy = OptionalStudyKeyStrategy()
+        filters = {"other": "param"}
+        key, new_filters = strategy.process(filters)
+
+        assert key is None
+        assert "studyKey" not in new_filters
+        assert new_filters["other"] == "param"