From d946dfcac0ba1502fd4014b1b816f8190a0084bf Mon Sep 17 00:00:00 2001 From: MobyNL Date: Tue, 18 Aug 2026 20:57:08 +0200 Subject: [PATCH] feat: simulate timeouts and truncated responses Blocking a request answers it, cleanly and at once. These do the opposite: they make a request fail the way a network fails, which is usually the least exercised path in an application under test. Simulate Timeout holds a request and then drops it, without contacting the server. That is what a client sees when a service accepts a connection and then says nothing, and it is a different test from a service that answers with an error - an application that handles a 504 correctly may still wait forever when nothing arrives at all. It is dropped after the hold so a client with a longer timeout than the hold does not wait indefinitely. Simulate Truncated Response cuts a response short while it still claims its full length. The mismatch is the fault: the body is replaced through raw_content because set_content would correct content-length and leave a merely shorter, perfectly valid response. A compressed body is cut in its compressed form, so what arrives is broken rather than a valid shorter document, which is the more realistic failure. A response with no body, and one already shorter than the cut, are left alone and say so rather than appearing to have done something. Connection reset is not a new keyword: Block Requests with mode=RESET already drops a connection, and a second way to do one thing is not worth adding to a surface about to be frozen. The failure simulation documentation points there. Bandwidth throttling stays unsupported, and the documentation says why rather than leaving it looking forgotten: mitmproxy hands a response body to a synchronous callback with no way to wait between chunks, so the only implementable version would delay the whole body and deliver it in one piece, which is Add Response Delay under a name promising more than it does. The request hook is now asynchronous, because a rule can wait in it. mitmproxy awaits an addon hook that returns a coroutine, and a held request does not stall the others - covered by a test that holds one request and requires another to pass through meanwhile, and by an acceptance case that does the same through a real proxy. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 12 ++ MitmLibrary/__init__.py | 102 +++++++++++++++ MitmLibrary/failures.py | 100 +++++++++++++++ MitmLibrary/interceptor.py | 11 +- MitmLibrary/rules.py | 4 +- README.md | 18 +++ atest/testcases/Test.robot | 29 +++++ tests/test_failures.py | 244 ++++++++++++++++++++++++++++++++++++ tests/test_interceptor.py | 20 +-- tests/test_request_rules.py | 40 +++--- 10 files changed, 549 insertions(+), 31 deletions(-) create mode 100644 MitmLibrary/failures.py create mode 100644 tests/test_failures.py diff --git a/CHANGELOG.md b/CHANGELOG.md index af882e9..56b2e7e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -82,6 +82,18 @@ Two behaviour changes come with it: than leaving the proxy to fail to start for an unstated reason. - `Start Mitm Proxy` takes `proxy_auth`, requiring clients to authenticate before the proxy serves them. +- Failure simulation, for the paths an application exercises least. `Simulate Timeout` + holds a request and then drops it without contacting the server, which is what a client + sees when a service accepts a connection and says nothing - a different test from a + service that answers with an error. `Simulate Truncated Response` cuts an answer short + while it still claims its full length, which is what a connection dropped mid-answer + looks like. A dropped connection is `Block Requests` with `mode=RESET` rather than a + keyword of its own. + + Bandwidth throttling is deliberately not supported: mitmproxy hands a response body to + a synchronous callback with no way to wait between chunks, so the only implementable + version would delay the whole body and deliver it in one piece, which is + `Add Response Delay` under a name that would promise more than it does. - `Get Proxy Rules` returns the loaded rules in the order they are applied. - Blocking rules have an alias, like every other rule, so they are removed the same way. - Rules survive a restart of the proxy: the registry outlives it, and only the addon diff --git a/MitmLibrary/__init__.py b/MitmLibrary/__init__.py index 82d70fd..43bdc56 100644 --- a/MitmLibrary/__init__.py +++ b/MitmLibrary/__init__.py @@ -19,6 +19,7 @@ from robot.api.deco import keyword, library, not_keyword from robot.utils import DotDict, timestr_to_secs +from MitmLibrary.failures import TimeoutAction, TruncateAction from MitmLibrary.interceptor import Interceptor from MitmLibrary.listener import LibraryListener from MitmLibrary.matching import ANY_METHOD, MatchMode, UrlMatcher @@ -153,6 +154,25 @@ class MitmLibrary: | # Send everything on through the network's own proxy | Start Mitm Proxy mode=upstream:http://corporate-proxy:3128 + = Failure simulation = + Most rules make a request succeed differently. These make it fail the way a network + fails, which is usually the least exercised path in an application: + + - `Simulate Timeout` holds a request and then drops it, so the client waits for an + answer that never comes. + - `Simulate Truncated Response` cuts an answer short while it still claims to be the + full length. + - A dropped connection is `Block Requests` with `mode=RESET`, which is where that + behaviour lives rather than in a keyword of its own. + + How a client reports any of these depends on the HTTP library it uses, so assert that + a request failed rather than on the particular error it raised. + + Bandwidth throttling is not supported. mitmproxy hands a response body to a + synchronous callback, with no way to wait between chunks, so the only implementable + version would delay the whole body and deliver it in one piece - which is + `Add Response Delay`, honestly named. + = Mitm Certificates = To test with SSL verification or use a browser without ignoring certificates, you need to set up certificates related to mitm. Follow the guide on the @@ -681,6 +701,88 @@ def redirect_requests_to_host( alias, url, method, match, times, RedirectAction(host, port, scheme) ) + @keyword + def simulate_timeout( + self, + alias: str, + url: str, + hold: str = "60s", + method: str = ANY_METHOD, + match: MatchMode = MatchMode.SUBSTRING, + times: int = 0, + ) -> None: + """Holds matching requests and then drops them, without contacting the server. + + This is what a client sees when a service accepts its connection and then says + nothing, and it is a different test from a service that answers with an error: an + application that handles a 504 correctly may still wait forever when nothing + arrives at all. + + Hold the request for longer than the client's own timeout, so the client is the + one that gives up. Other traffic is unaffected while a request is held. + + - `alias`: The handle for this rule. Reusing an alias replaces the rule that + already uses it. + - `url`: The pattern the request url is compared against. See `match`. + - `hold`: How long to hold the request, in Robot Framework time format. It is + then dropped, so the client does not wait forever if its own timeout is longer. + - `method`: Only match this HTTP method. `ANY` matches every method. + - `match`: How `url` is interpreted. See the `Matching` section. + - `times`: How often the rule may be applied. `0` means unlimited. + + Example: + | Simulate Timeout hang /api/orders hold=30s + """ + self._require_proxy() + self._add_rule( + alias, url, method, match, times, TimeoutAction(timestr_to_secs(hold), hold) + ) + + @keyword + def simulate_truncated_response( + self, + alias: str, + url: str, + keep_bytes: Optional[int] = None, + keep_fraction: float = 0.5, + method: str = ANY_METHOD, + match: MatchMode = MatchMode.SUBSTRING, + times: int = 0, + ) -> None: + """Cuts matching responses short while they still claim to be the full length. + + The response keeps saying how long its body was meant to be, so a client reads + what arrived, waits for the rest and eventually gives up. That is what a + connection dropped mid-answer looks like, and it is the case that finds parsers + which assume a body is either complete or absent. + + How the client reports it differs per HTTP library, so assert that the request + failed rather than on a particular error. + + - `alias`: The handle for this rule. Reusing an alias replaces the rule that + already uses it. + - `url`: The pattern the request url is compared against. See `match`. + - `keep_bytes`: How many bytes of the body to keep. Overrides `keep_fraction`. + - `keep_fraction`: How much of the body to keep, as a fraction. Half by default. + - `method`: Only match this HTTP method. `ANY` matches every method. + - `match`: How `url` is interpreted. See the `Matching` section. + - `times`: How often the rule may be applied. `0` means unlimited. + + A response with no body, and one already shorter than what would be kept, are + left alone and say so in the log. + + A compressed body is cut in its compressed form, so what arrives is not decodable + rather than being a valid shorter document. That is the more realistic failure, + and the more interesting one to test against. + + Example: + | Simulate Truncated Response cut /api/orders keep_bytes=10 + """ + self._require_proxy() + self._add_rule( + alias, url, method, match, times, TruncateAction(keep_bytes, keep_fraction) + ) + @keyword def remove_rule(self, alias: str) -> None: """Removes the rule with the given alias. diff --git a/MitmLibrary/failures.py b/MitmLibrary/failures.py new file mode 100644 index 0000000..f350c13 --- /dev/null +++ b/MitmLibrary/failures.py @@ -0,0 +1,100 @@ +""" +This file defines the rules that break traffic on purpose. + +Blocking a request answers it, cleanly and immediately. These do the opposite: they make a +request fail the way a real network fails, so a suite can see what the application under +test does when a service hangs or an answer arrives half-finished. Those paths are usually +the least exercised and the most likely to be wrong. + +They live apart from the other rules because they are the most version-sensitive part of +the library: each one depends on a mitmproxy behaviour that is documented but not promised, +and if one has to be reverted it should be possible to do that without touching the model +everything else is built on. +""" + +import asyncio +from dataclasses import dataclass +from typing import Any, Dict, Optional + +from mitmproxy import http +from robot.api import logger + +from MitmLibrary.rules import Action, kill_flow, Phase, Priority + + +@dataclass(frozen=True) +class TimeoutAction(Action): + """Holds a request and then drops it, so the client runs into its own timeout. + + The request is never sent, so the service is not involved at all: this is what a + client sees when a service accepts a connection and then says nothing. Holding rather + than answering is the point - an application that handles a 504 correctly may still + hang forever when nothing arrives. + """ + + hold_seconds: float = 60.0 + hold: str = "" + + phase = Phase.REQUEST + priority = Priority.TERMINAL + + def apply(self, flow: http.HTTPFlow) -> bool: # pragma: no cover - async path is used + raise NotImplementedError("A timeout can only be applied from the async hook.") + + async def apply_async(self, flow: http.HTTPFlow) -> bool: + await asyncio.sleep(self.hold_seconds) + kill_flow(flow) + return True + + def describe(self) -> Dict[str, Any]: + return {"type": "timeout", "hold": self.hold, "hold_seconds": self.hold_seconds} + + +@dataclass(frozen=True) +class TruncateAction(Action): + """Cuts a response short while it still claims to be the full length. + + The `content-length` header keeps saying how long the body was meant to be, so a + client reads it, waits for the rest, and eventually gives up. That mismatch is the + fault being injected, which is why the body is replaced through `raw_content`: + `set_content` would helpfully correct the header and there would be nothing wrong + with the response at all. + """ + + keep_bytes: Optional[int] = None + keep_fraction: float = 0.5 + + phase = Phase.RESPONSE + priority = Priority.MUTATE + + def apply(self, flow: http.HTTPFlow) -> bool: + if flow.response is None: + return False + body = flow.response.raw_content + if body is None: + # A streamed response has no body to cut here, and a 204 or 304 has none at + # all. Saying so beats a rule that silently did nothing. + logger.info("There was no body to truncate, so the response is unchanged.") + return False + keep = self._keep(len(body)) + if keep >= len(body): + logger.info( + f"The body is {len(body)} bytes, which is not longer than the " + f"{keep} bytes to keep, so the response is unchanged." + ) + return False + flow.response.raw_content = body[:keep] + return False + + def _keep(self, length: int) -> int: + """How many bytes to keep, from either a count or a fraction of the body.""" + if self.keep_bytes is not None: + return max(0, self.keep_bytes) + return max(0, int(length * self.keep_fraction)) + + def describe(self) -> Dict[str, Any]: + return { + "type": "truncate", + "keep_bytes": self.keep_bytes, + "keep_fraction": self.keep_fraction, + } diff --git a/MitmLibrary/interceptor.py b/MitmLibrary/interceptor.py index 79ad210..30f4a99 100644 --- a/MitmLibrary/interceptor.py +++ b/MitmLibrary/interceptor.py @@ -28,15 +28,20 @@ def set_console_logging(self, value: bool) -> None: """Enables or disables reporting each manipulation on the console.""" self.log_to_console = value - def request(self, flow: http.HTTPFlow) -> None: - """Applies the rules that act before the request is sent.""" + async def request(self, flow: http.HTTPFlow) -> None: + """Applies the rules that act before the request is sent. + + Asynchronous because a rule may wait here: a simulated timeout holds the request + rather than answering it. mitmproxy awaits an addon hook that returns a + coroutine, and waiting in one flow does not hold up the others. + """ for rule in self.registry.snapshot(Phase.REQUEST): if not self._applies(rule, flow): continue if not self.registry.consume(rule): continue self._log(rule, flow) - if rule.action.apply(flow): + if await rule.action.apply_async(flow): return async def response(self, flow: http.HTTPFlow) -> None: diff --git a/MitmLibrary/rules.py b/MitmLibrary/rules.py index dcfaf90..369ead2 100644 --- a/MitmLibrary/rules.py +++ b/MitmLibrary/rules.py @@ -107,7 +107,7 @@ class BlockAction(Action): def apply(self, flow: http.HTTPFlow) -> bool: if self.mode is BlockMode.RESET: - _kill(flow) + kill_flow(flow) return True flow.response = http.Response.make( self.status_code, safe_str(self.body) if self.body is not None else b"" @@ -351,7 +351,7 @@ def _is_default_port(scheme: str, port: int) -> bool: return (scheme == "http" and port == 80) or (scheme == "https" and port == 443) -def _kill(flow: http.HTTPFlow) -> None: +def kill_flow(flow: http.HTTPFlow) -> None: """Drops the connection, if mitmproxy still lets us. `kill()` raises when the flow is no longer killable, which happens when something else diff --git a/README.md b/README.md index 7395994..ddccfa9 100644 --- a/README.md +++ b/README.md @@ -78,6 +78,24 @@ All matching rules are applied. A rule that blocks a request ends it and nothing runs; otherwise `Set Response` runs before rules that change part of a response, which run before delays, so combinations behave predictably rather than overwriting each other. +### Simulating failures + +Most rules make a request succeed differently. These make it fail the way a network does: + +```robotframework +Simulate Timeout hang /api/orders hold=30s +Simulate Truncated Response cut /api/orders keep_bytes=10 +Block Requests drop /api/orders mode=RESET +``` + +How a client reports any of these depends on the HTTP library it uses, so assert that the +request failed rather than on the particular error. + +Bandwidth throttling is not supported: mitmproxy hands a response body to a synchronous +callback with no way to wait between chunks, so the only implementable version would delay +the whole body and deliver it in one piece — which is what `Add Response Delay` already +does, honestly named. + ### Recording The proxy can also remember what went through it, so a suite can assert on what the diff --git a/atest/testcases/Test.robot b/atest/testcases/Test.robot index df50b84..5c0175b 100644 --- a/atest/testcases/Test.robot +++ b/atest/testcases/Test.robot @@ -231,6 +231,35 @@ Waiting For A Request That Never Comes Fails Recording Keywords Explain Themselves When Recording Is Off Run Keyword And Expect Error *Start Recording* Get Recorded Requests +A Held Request Runs Into The Client's Own Timeout + [Documentation] How a client reports this differs per HTTP library, so this asserts + ... that the request failed rather than on a particular error. The hold is far + ... longer than the client's timeout, so the client is the one that gives up and + ... the test does not depend on the two being close together. + Simulate Timeout hang /test_post hold=30s + ${start} Get Time epoch + Run Keyword And Expect Error * + ... POST On Session alias=proxy url=test_post/1 timeout=2 + ${end} Get Time epoch + # It failed because it waited, not because it was refused outright. + Should Be True ${end} - ${start} >= 2 The request failed without waiting + +Other Traffic Keeps Flowing While A Request Is Held + [Documentation] Holding must not stall the proxy for everything else. + Simulate Timeout hang /test_get hold=30s + Check POST Response smaller than 2 ${200} + +A Truncated Response Fails The Client + [Documentation] The response still claims its full length, so the client waits for + ... a rest that never comes. Which error that produces is the client's business. + Simulate Truncated Response cut /test_post keep_bytes=5 + Run Keyword And Expect Error * + ... POST On Session alias=proxy url=test_post/1 timeout=5 + +A Response Shorter Than The Cut Is Left Alone + Simulate Truncated Response cut /test_post keep_bytes=1000 + Check POST Response smaller than 2 ${200} + Delayed Response With Post Check POST Response smaller than 2 ${200} Add Response Delay alias=delay url=test_post delay=5s diff --git a/tests/test_failures.py b/tests/test_failures.py new file mode 100644 index 0000000..667e482 --- /dev/null +++ b/tests/test_failures.py @@ -0,0 +1,244 @@ +"""Tests for the rules that break traffic on purpose. + +These are the most version-sensitive rules in the library, because each depends on a +mitmproxy behaviour that is documented but not promised. The assertions are written to +notice if one of those behaviours changes, rather than only that the keyword ran. +""" + +import asyncio +import unittest +from unittest.mock import patch + +from mitmproxy import exceptions, http +from mitmproxy.test import tflow, tutils +from robot.api import logger + +from MitmLibrary.failures import TimeoutAction, TruncateAction +from MitmLibrary.interceptor import Interceptor +from MitmLibrary.matching import MatchMode, UrlMatcher +from MitmLibrary.rules import Rule, RuleRegistry + + +def make_flow(url="http://example.com/api/users", body=b"a full response body"): + """A real flow whose content-length agrees with its body, as a server's would.""" + flow = tflow.tflow(req=tutils.treq(), resp=tutils.tresp(content=body)) + flow.request.url = url + flow.response.headers["content-length"] = str(len(body)) + return flow + + +class FailureTestCase(unittest.TestCase): + def setUp(self): + self.registry = RuleRegistry() + self.interceptor = Interceptor(self.registry, log_to_console=False) + + def add(self, alias, action, url="/api", method="ANY", times=0): + self.registry.add( + Rule(alias, UrlMatcher(url, MatchMode.SUBSTRING, method), action, times) + ) + + def send(self, flow): + asyncio.run(self.interceptor.request(flow)) + + def respond(self, flow): + asyncio.run(self.interceptor.response(flow)) + + +class TestTimeout(unittest.TestCase): + """The wait is asserted with a patched sleep; waiting for real proves nothing extra.""" + + def setUp(self): + self.registry = RuleRegistry() + self.interceptor = Interceptor(self.registry, log_to_console=False) + self.registry.add( + Rule("hang", UrlMatcher("/api"), TimeoutAction(30.0, "30s")) + ) + + def test_the_request_is_held_for_the_configured_time(self): + flow = make_flow() + waited = [] + + async def record(seconds): + waited.append(seconds) + + with patch("MitmLibrary.failures.asyncio.sleep", record): + asyncio.run(self.interceptor.request(flow)) + self.assertEqual(waited, [30.0]) + + def test_the_request_is_dropped_afterwards(self): + """Without this the client waits forever when its own timeout is longer.""" + flow = make_flow() + flow.kill = lambda: setattr(flow, "killed", True) + + async def record(_seconds): + self.assertFalse(getattr(flow, "killed", False)) # held, not yet dropped + + with patch("MitmLibrary.failures.asyncio.sleep", record): + asyncio.run(self.interceptor.request(flow)) + self.assertTrue(flow.killed) + + def test_the_request_never_reaches_the_server(self): + """A timeout is not a slow answer: the request is not sent at all.""" + flow = make_flow() + + async def record(_seconds): + return None + + with patch("MitmLibrary.failures.asyncio.sleep", record): + asyncio.run(self.interceptor.request(flow)) + self.assertIsNotNone(flow.error) + + def test_a_timeout_ends_the_flow(self): + """Nothing after it should run, since there is no longer a request to change.""" + later = Rule("later", UrlMatcher("/api"), TimeoutAction(1.0, "1s")) + self.registry.add(later) + + async def record(_seconds): + return None + + with patch("MitmLibrary.failures.asyncio.sleep", record): + asyncio.run(self.interceptor.request(make_flow())) + self.assertEqual(later.used, 0) + + def test_holding_one_request_does_not_hold_up_another(self): + """The whole point of holding rather than blocking: other traffic keeps flowing.""" + self.registry.clear() + self.registry.add(Rule("hang", UrlMatcher("/slow"), TimeoutAction(0.4, "0.4s"))) + slow = make_flow(url="http://example.com/slow") + fast = make_flow() + + async def scenario(): + held = asyncio.create_task(self.interceptor.request(slow)) + await asyncio.sleep(0) + start = asyncio.get_running_loop().time() + await self.interceptor.request(fast) + elapsed = asyncio.get_running_loop().time() - start + await held + return elapsed + + self.assertLess(asyncio.run(scenario()), 0.2) + + def test_a_flow_that_cannot_be_killed_does_not_raise(self): + flow = make_flow() + flow.kill = lambda: (_ for _ in ()).throw( + exceptions.ControlException("Flow is not killable.") + ) + flow.live = False + + async def record(_seconds): + return None + + with patch("MitmLibrary.failures.asyncio.sleep", record): + asyncio.run(self.interceptor.request(flow)) # must not raise + + def test_the_rule_reports_what_it_was_asked_for(self): + described = self.registry.describe()[0] + self.assertEqual(described.type, "timeout") + self.assertEqual(described.hold, "30s") + self.assertEqual(described.hold_seconds, 30.0) + self.assertEqual(described.phase, "request") + + +class TestTruncate(FailureTestCase): + def test_the_body_is_cut_short(self): + self.add("cut", TruncateAction(keep_bytes=6)) + flow = make_flow(body=b"a full response body") + self.respond(flow) + self.assertEqual(flow.response.raw_content, b"a full") + + def test_the_declared_length_is_left_alone(self): + """The mismatch is the fault being injected. + + set_content would correct content-length and leave a perfectly valid, merely + shorter response, which is not a failure at all. This assertion is what would + notice a mitmproxy version that started normalising the header on write. + """ + self.add("cut", TruncateAction(keep_bytes=6)) + flow = make_flow(body=b"a full response body") + self.respond(flow) + self.assertEqual(flow.response.headers["content-length"], "20") + self.assertEqual(len(flow.response.raw_content), 6) + + def test_half_the_body_is_kept_by_default(self): + self.add("cut", TruncateAction()) + flow = make_flow(body=b"0123456789") + self.respond(flow) + self.assertEqual(flow.response.raw_content, b"01234") + + def test_a_fraction_can_be_given(self): + self.add("cut", TruncateAction(keep_fraction=0.25)) + flow = make_flow(body=b"0123456789") + self.respond(flow) + self.assertEqual(flow.response.raw_content, b"01") + + def test_a_byte_count_wins_over_a_fraction(self): + self.add("cut", TruncateAction(keep_bytes=3, keep_fraction=0.9)) + flow = make_flow(body=b"0123456789") + self.respond(flow) + self.assertEqual(flow.response.raw_content, b"012") + + def test_nothing_at_all_can_be_kept(self): + self.add("cut", TruncateAction(keep_bytes=0)) + flow = make_flow(body=b"0123456789") + self.respond(flow) + self.assertEqual(flow.response.raw_content, b"") + self.assertEqual(flow.response.headers["content-length"], "10") + + def test_a_negative_count_keeps_nothing_rather_than_failing(self): + self.add("cut", TruncateAction(keep_bytes=-5)) + flow = make_flow(body=b"0123456789") + self.respond(flow) + self.assertEqual(flow.response.raw_content, b"") + + def test_a_body_that_is_already_short_enough_is_left_alone(self): + self.add("cut", TruncateAction(keep_bytes=100)) + flow = make_flow(body=b"short") + with patch.object(logger, "info") as mock_info: + self.respond(flow) + self.assertEqual(flow.response.raw_content, b"short") + logged = " ".join(str(call.args[0]) for call in mock_info.call_args_list) + self.assertIn("unchanged", logged) + + def test_a_response_without_a_body_says_so(self): + """A 204 or a streamed response has nothing here to cut.""" + self.add("cut", TruncateAction(keep_bytes=1)) + flow = make_flow() + flow.response = http.Response.make(204) + flow.response.raw_content = None + with patch.object(logger, "info") as mock_info: + self.respond(flow) + logged = " ".join(str(call.args[0]) for call in mock_info.call_args_list) + self.assertIn("no body to truncate", logged) + + def test_a_flow_without_a_response_is_a_no_op(self): + self.add("cut", TruncateAction(keep_bytes=1)) + flow = make_flow() + flow.response = None + self.respond(flow) # must not raise + + def test_a_compressed_body_is_cut_in_its_compressed_form(self): + """What arrives is not a valid shorter document but a broken one. + + That is the more realistic failure and the more interesting one to test an + application with. mitmproxy does not raise on a body it cannot decompress - it + hands back nothing - so the assertion is that the original text is gone, not that + decoding fails. + """ + self.add("cut", TruncateAction(keep_fraction=0.5)) + flow = make_flow(body=b"a full response body repeated over and over and over") + flow.response.encode("gzip") + compressed = flow.response.raw_content + self.respond(flow) + self.assertEqual(len(flow.response.raw_content), len(compressed) // 2) + self.assertNotIn(b"a full response body", flow.response.get_content()) + + def test_the_rule_reports_what_it_was_asked_for(self): + self.add("cut", TruncateAction(keep_bytes=6)) + described = self.registry.describe()[0] + self.assertEqual(described.type, "truncate") + self.assertEqual(described.keep_bytes, 6) + self.assertEqual(described.phase, "response") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_interceptor.py b/tests/test_interceptor.py index f466270..4b0bbb6 100644 --- a/tests/test_interceptor.py +++ b/tests/test_interceptor.py @@ -59,26 +59,30 @@ def respond(self, flow): """Runs the async response hook, as mitmproxy would.""" asyncio.run(self.interceptor.response(flow)) + def send(self, flow): + """Runs the async request hook, as mitmproxy would.""" + asyncio.run(self.interceptor.request(flow)) + class TestBlocking(InterceptorTestCase): def test_respond_mode_answers_without_reaching_the_server(self): self.add("block", BlockAction(BlockMode.RESPOND, 403)) flow = make_flow() - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.response.status_code, 403) flow.kill.assert_not_called() def test_respond_mode_can_carry_a_body(self): self.add("block", BlockAction(BlockMode.RESPOND, 503, "maintenance")) flow = make_flow() - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.response.content, b"maintenance") self.assertEqual(flow.response.status_code, 503) def test_reset_mode_drops_the_connection(self): self.add("block", BlockAction(BlockMode.RESET)) flow = make_flow() - self.interceptor.request(flow) + self.send(flow) flow.kill.assert_called_once() def test_reset_mode_does_not_raise_on_a_flow_that_cannot_be_killed(self): @@ -87,13 +91,13 @@ def test_reset_mode_does_not_raise_on_a_flow_that_cannot_be_killed(self): flow = make_flow() flow.killable = False flow.kill.side_effect = exceptions.ControlException("Flow is not killable.") - self.interceptor.request(flow) # must not raise + self.send(flow) # must not raise flow.kill.assert_not_called() def test_a_request_that_does_not_match_is_left_alone(self): self.add("block", BlockAction(), url="/orders") flow = make_flow() - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.response.status_code, 200) def test_blocking_ends_the_flow(self): @@ -101,7 +105,7 @@ def test_blocking_ends_the_flow(self): self.add("first", BlockAction(BlockMode.RESPOND, 403)) second = self.add("second", BlockAction(BlockMode.RESPOND, 503)) flow = make_flow() - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.response.status_code, 403) self.assertEqual(second.used, 0) @@ -250,9 +254,9 @@ def test_an_exhausted_rule_is_skipped_in_the_response_hook(self): def test_an_exhausted_rule_is_skipped_in_the_request_hook(self): rule = self.add("block", BlockAction(BlockMode.RESPOND, 403), times=1) - self.interceptor.request(make_flow()) + self.send(make_flow()) flow = make_flow() - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.response.status_code, 200) self.assertEqual(rule.used, 1) diff --git a/tests/test_request_rules.py b/tests/test_request_rules.py index c79c158..f62a069 100644 --- a/tests/test_request_rules.py +++ b/tests/test_request_rules.py @@ -50,36 +50,40 @@ def add(self, alias, action, url="/api", method="ANY", times=0): def respond(self, flow): asyncio.run(self.interceptor.response(flow)) + def send(self, flow): + """Runs the async request hook, as mitmproxy would.""" + asyncio.run(self.interceptor.request(flow)) + class TestRequestHeaders(RequestRuleTestCase): def test_a_header_can_be_added(self): self.add("auth", RequestHeadersAction({"Authorization": "Bearer token"})) flow = make_flow() - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.request.headers["Authorization"], "Bearer token") def test_an_existing_header_is_replaced(self): self.add("auth", RequestHeadersAction({"Authorization": "Bearer new"})) flow = make_flow(headers={"Authorization": "Bearer old"}) - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.request.headers["Authorization"], "Bearer new") def test_other_headers_are_left_alone(self): """Merging is the point: setting one header must not drop the rest.""" self.add("auth", RequestHeadersAction({"Authorization": "Bearer token"})) flow = make_flow(headers={"X-Kept": "yes"}) - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.request.headers["X-Kept"], "yes") def test_a_header_can_be_removed(self): self.add("cookies", RequestHeadersAction(None, ["Cookie"])) flow = make_flow(headers={"Cookie": "session=1"}) - self.interceptor.request(flow) + self.send(flow) self.assertNotIn("Cookie", flow.request.headers) def test_removing_a_header_that_is_not_there_is_a_no_op(self): self.add("cookies", RequestHeadersAction(None, ["Cookie"])) - self.interceptor.request(make_flow()) # must not raise + self.send(make_flow()) # must not raise def test_removing_and_setting_the_same_header_leaves_one_value(self): """Headers can repeat; naming one in both is how a suite collapses them.""" @@ -87,7 +91,7 @@ def test_removing_and_setting_the_same_header_leaves_one_value(self): flow = make_flow() flow.request.headers.add("Accept", "text/plain") flow.request.headers.add("Accept", "text/html") - self.interceptor.request(flow) + self.send(flow) self.assertEqual( flow.request.headers.get_all("Accept"), ["application/json"] ) @@ -97,14 +101,14 @@ class TestRequestBody(RequestRuleTestCase): def test_the_body_is_replaced(self): self.add("payload", RequestBodyAction('{"id": 1}')) flow = make_flow(method="POST", body=b"original request") - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.request.content, b'{"id": 1}') def test_the_content_length_is_updated(self): """A declared length that no longer matches makes the request unreadable.""" self.add("payload", RequestBodyAction("short")) flow = make_flow(method="POST", body=b"a much longer original body") - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.request.headers["content-length"], "5") @@ -112,7 +116,7 @@ class TestRewrite(RequestRuleTestCase): def test_the_whole_url_is_replaced(self): self.add("v2", RewriteAction("https://other.example.com:8443/api/v2/users?x=1")) flow = make_flow() - self.interceptor.request(flow) + self.send(flow) self.assertEqual( flow.request.pretty_url, "https://other.example.com:8443/api/v2/users?x=1" ) @@ -120,7 +124,7 @@ def test_the_whole_url_is_replaced(self): def test_the_parts_of_the_request_stay_consistent(self): self.add("v2", RewriteAction("https://other.example.com:8443/api/v2/users")) flow = make_flow() - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.request.scheme, "https") self.assertEqual(flow.request.host, "other.example.com") self.assertEqual(flow.request.port, 8443) @@ -133,7 +137,7 @@ def test_the_host_header_follows_the_url(self): """ self.add("v2", RewriteAction("https://other.example.com:8443/api/v2/users")) flow = make_flow(headers={"Host": "example.com"}) - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.request.host_header, "other.example.com:8443") @@ -141,7 +145,7 @@ class TestRedirect(RequestRuleTestCase): def test_the_host_is_replaced_and_the_path_is_kept(self): self.add("stub", RedirectAction("127.0.0.1", 8000)) flow = make_flow(url="http://api.example.com/api/users?page=2") - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.request.host, "127.0.0.1") self.assertEqual(flow.request.port, 8000) self.assertEqual(flow.request.path, "/api/users?page=2") @@ -149,13 +153,13 @@ def test_the_host_is_replaced_and_the_path_is_kept(self): def test_the_original_port_is_kept_when_none_is_given(self): self.add("stub", RedirectAction("127.0.0.1")) flow = make_flow(url="http://api.example.com:9000/api/users") - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.request.port, 9000) def test_the_scheme_can_be_changed(self): self.add("stub", RedirectAction("127.0.0.1", 8000, "https")) flow = make_flow(url="http://api.example.com/api/users") - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.request.scheme, "https") def test_the_host_header_is_updated(self): @@ -163,7 +167,7 @@ def test_the_host_header_is_updated(self): self.add("stub", RedirectAction("127.0.0.1", 8000)) flow = make_flow(url="http://api.example.com/api/users", headers={"Host": "api.example.com"}) - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.request.host_header, "127.0.0.1:8000") def test_a_default_port_is_left_out_of_the_host_header(self): @@ -171,14 +175,14 @@ def test_a_default_port_is_left_out_of_the_host_header(self): self.add("stub", RedirectAction("other.example.com", 80)) flow = make_flow(url="http://api.example.com/api/users", headers={"Host": "api.example.com"}) - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.request.host_header, "other.example.com") def test_a_default_https_port_is_left_out_too(self): self.add("stub", RedirectAction("other.example.com", 443, "https")) flow = make_flow(url="http://api.example.com/api/users", headers={"Host": "api.example.com"}) - self.interceptor.request(flow) + self.send(flow) self.assertEqual(flow.request.host_header, "other.example.com") def test_a_request_without_a_host_header_does_not_gain_one(self): @@ -186,7 +190,7 @@ def test_a_request_without_a_host_header_does_not_gain_one(self): self.add("stub", RedirectAction("127.0.0.1", 8000)) flow = make_flow() self.assertIsNone(flow.request.host_header) # tflow builds one without it - self.interceptor.request(flow) + self.send(flow) self.assertIsNone(flow.request.host_header)