Add type hints and a test

Author: Gallaecio

scrapy_zyte_api/_session.py

@@ -982,23 +982,27 @@ class LocationSessionConfig(SessionConfig):
     as a parameter.
     """
 
-    def params(self, request):
+    def params(self, request: Request) -> Dict[str, Any]:
         if not (location := self.location(request)):
             return super().params(request)
         return self.location_params(request, location)
 
-    def check(self, response, request):
+    def check(self, response: Response, request: Request) -> bool:
         if not (location := self.location(request)):
             return super().check(response, request)
         return self.location_check(response, request, location)
 
-    def location_params(self, request, location):
+    def location_params(
+        self, request: Request, location: Dict[str, Any]
+    ) -> Dict[str, Any]:
         """Like :class:`SessionConfig.params
         <scrapy_zyte_api.SessionConfig.params>`, but it is only called when a
         location it set, and gets that *location* as a parameter."""
         return super().params(request)
 
-    def location_check(self, response, request, location):
+    def location_check(
+        self, response: Response, request: Request, location: Dict[str, Any]
+    ) -> bool:
         """Like :class:`SessionConfig.check
         <scrapy_zyte_api.SessionConfig.check>`, but it is only called when a
         location it set, and gets that *location* as a parameter."""

tests/test_sessions.py

@@ -17,6 +17,7 @@
 from scrapy_zyte_api import (
     SESSION_AGGRESSIVE_RETRY_POLICY,
     SESSION_DEFAULT_RETRY_POLICY,
+    LocationSessionConfig,
     SessionConfig,
     is_session_init_request,
     session_config,
@@ -2080,6 +2081,119 @@ class CustomSessionConfig(SessionConfig):
             pass
 
 
+@ensureDeferred
+async def test_location_session_config(mockserver):
+    pytest.importorskip("web_poet")
+
+    @session_config(
+        [
+            "postal-code-10001.example",
+            "postal-code-10001-fail.example",
+            "postal-code-10001-alternative.example",
+        ]
+    )
+    class CustomSessionConfig(LocationSessionConfig):
+
+        def location_params(
+            self, request: Request, location: Dict[str, Any]
+        ) -> Dict[str, Any]:
+            assert location == {"postalCode": "10002"}
+            return {
+                "actions": [
+                    {
+                        "action": "setLocation",
+                        "address": {"postalCode": "10001"},
+                    }
+                ]
+            }
+
+        def location_check(
+            self, response: Response, request: Request, location: Dict[str, Any]
+        ) -> bool:
+            assert location == {"postalCode": "10002"}
+            domain = urlparse_cached(request).netloc
+            return "fail" not in domain
+
+        def pool(self, request: Request) -> str:
+            domain = urlparse_cached(request).netloc
+            if domain == "postal-code-10001-alternative.example":
+                return "postal-code-10001.example"
+            return domain
+
+    settings = {
+        "RETRY_TIMES": 0,
+        "ZYTE_API_URL": mockserver.urljoin("/"),
+        "ZYTE_API_SESSION_ENABLED": True,
+        # We set a location to force the location-specific methods of the
+        # session config class to be called, but we set the wrong location so
+        # that the test would not pass were it not for our custom
+        # implementation which ignores the input location and instead sets the
+        # right one.
+        "ZYTE_API_SESSION_LOCATION": {"postalCode": "10002"},
+        "ZYTE_API_SESSION_MAX_BAD_INITS": 1,
+    }
+
+    class TestSpider(Spider):
+        name = "test"
+        start_urls = [
+            "https://postal-code-10001.example",
+            "https://postal-code-10001-alternative.example",
+            "https://postal-code-10001-fail.example",
+        ]
+
+        def start_requests(self):
+            for url in self.start_urls:
+                yield Request(
+                    url,
+                    meta={
+                        "zyte_api_automap": {
+                            "actions": [
+                                {
+                                    "action": "setLocation",
+                                    "address": {"postalCode": "10001"},
+                                }
+                            ]
+                        },
+                    },
+                )
+
+        def parse(self, response):
+            pass
+
+    crawler = await get_crawler(settings, spider_cls=TestSpider, setup_engine=False)
+    await crawler.crawl()
+
+    session_stats = {
+        k: v
+        for k, v in crawler.stats.get_stats().items()
+        if k.startswith("scrapy-zyte-api/sessions")
+    }
+    assert session_stats == {
+        "scrapy-zyte-api/sessions/pools/postal-code-10001.example/init/check-passed": 2,
+        "scrapy-zyte-api/sessions/pools/postal-code-10001.example/use/check-passed": 2,
+        "scrapy-zyte-api/sessions/pools/postal-code-10001-fail.example/init/check-failed": 1,
+    }
+
+    # Clean up the session config registry, and check it, otherwise we could
+    # affect other tests.
+
+    session_config_registry.__init__()  # type: ignore[misc]
+
+    crawler = await get_crawler(settings, spider_cls=TestSpider, setup_engine=False)
+    await crawler.crawl()
+
+    session_stats = {
+        k: v
+        for k, v in crawler.stats.get_stats().items()
+        if k.startswith("scrapy-zyte-api/sessions")
+    }
+    assert session_stats == {
+        "scrapy-zyte-api/sessions/pools/postal-code-10001.example/init/failed": 1,
+        "scrapy-zyte-api/sessions/pools/postal-code-10001-alternative.example/init/failed": 1,
+        "scrapy-zyte-api/sessions/pools/postal-code-10001-fail.example/init/failed": 1,
+    }
+
+
 @ensureDeferred
 async def test_session_refresh(mockserver):
     """If a response does not pass a session validity check, the session is