diff --git a/bin/tests/system/resend_loop/ans3/ans.py b/bin/tests/system/resend_loop/ans3/ans.py index 9dcdec83b5..37ed3482f0 100644 --- a/bin/tests/system/resend_loop/ans3/ans.py +++ b/bin/tests/system/resend_loop/ans3/ans.py @@ -28,16 +28,6 @@ from isctest.asyncserver import ( ) -def _get_cookie(qctx: QueryContext) -> dns.edns.CookieOption | None: - for o in qctx.query.options: - if o.otype == dns.edns.OptionType.COOKIE: - cookie = o - cookie.server = b"\x11\x22\x33\x44\x55\x66\x77\x88" - return cookie - - return None - - def rrset( qname: dns.name.Name | str, rtype: dns.rdatatype.RdataType, @@ -64,10 +54,19 @@ class ExampleNSHandler(QnameQtypeHandler, StaticResponseHandler): class CookieHandler(DomainHandler): domains = ["example."] + def _get_cookie(self, qctx: QueryContext) -> dns.edns.CookieOption | None: + for o in qctx.query.options: + if o.otype == dns.edns.OptionType.COOKIE: + cookie = o + cookie.server = b"\x11\x22\x33\x44\x55\x66\x77\x88" + return cookie + + return None + async def get_responses( self, qctx: QueryContext ) -> AsyncGenerator[DnsResponseSend, None]: - if cookie := _get_cookie(qctx): + if cookie := self._get_cookie(qctx): # If there is a client cookie, mock BADCOOKIE to trigger # the resend loop logic. qctx.response.use_edns(options=[cookie])