mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-22 15:37:45 +00:00
[sans-io] playbook: fix multiple placeholders of same type
This commit is contained in:
Maximilian Hils
parent
9251f7820a
commit
553f102d6e
@ -3,7 +3,7 @@ from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
|
||||
from mitmproxy.proxy2 import events, commands, layer
|
||||
from mitmproxy.proxy2 import commands, events, layer
|
||||
from . import tutils
|
||||
|
||||
|
||||
@ -102,63 +102,6 @@ def test_placeholder_type_mismatch(tplaybook):
|
||||
)
|
||||
|
||||
|
||||
def test_fork(tplaybook):
|
||||
"""Playbooks can be forked to test multiple execution streams."""
|
||||
assert (
|
||||
tplaybook
|
||||
>> TEvent()
|
||||
<< TCommand()
|
||||
)
|
||||
p2 = tplaybook.fork()
|
||||
p3 = tplaybook.fork()
|
||||
assert (
|
||||
tplaybook
|
||||
>> TEvent()
|
||||
<< TCommand()
|
||||
)
|
||||
assert (
|
||||
p2
|
||||
>> TEvent()
|
||||
<< TCommand()
|
||||
)
|
||||
assert len(tplaybook.actual) == len(tplaybook.expected) == 4
|
||||
assert len(p2.actual) == len(p2.expected) == 4
|
||||
assert len(p3.actual) == len(p3.expected) == 2
|
||||
|
||||
|
||||
def test_fork_placeholder(tplaybook):
|
||||
"""Forks require new placeholders."""
|
||||
f = tutils.Placeholder()
|
||||
flow = object()
|
||||
assert (
|
||||
tplaybook
|
||||
>> TEvent([flow])
|
||||
<< TCommand(f)
|
||||
)
|
||||
assert f() == flow
|
||||
p2 = tplaybook.fork()
|
||||
|
||||
p2_flow = p2.expected[0].commands[0]
|
||||
assert p2_flow != flow
|
||||
|
||||
# As we have forked, we need a new placeholder.
|
||||
f2 = tutils.Placeholder()
|
||||
assert (
|
||||
p2
|
||||
>> TEvent([p2_flow])
|
||||
<< TCommand(f2)
|
||||
)
|
||||
assert f2() == p2_flow
|
||||
|
||||
# re-using the old placeholder does not work.
|
||||
with pytest.raises(AssertionError, match="Playbook mismatch"):
|
||||
assert (
|
||||
p2
|
||||
>> TEvent([p2_flow])
|
||||
<< TCommand(f)
|
||||
)
|
||||
|
||||
|
||||
def test_unfinished(tplaybook):
|
||||
"""We show a warning when playbooks aren't asserted."""
|
||||
tplaybook >> TEvent()
|
||||
@ -209,3 +152,28 @@ def test_eq_placeholder():
|
||||
|
||||
b.foo._obj = 44
|
||||
assert not tutils._eq(a, b)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("swap", [False, True])
|
||||
def test_command_multiple_replies(tplaybook, swap):
|
||||
a = tutils.Placeholder(int)
|
||||
b = tutils.Placeholder(int)
|
||||
|
||||
command1 = TCommand(a)
|
||||
command2 = TCommand(b)
|
||||
|
||||
(tplaybook
|
||||
>> TEvent([1])
|
||||
<< command1
|
||||
>> TEvent([2])
|
||||
<< command2
|
||||
)
|
||||
if swap:
|
||||
tplaybook >> tutils.reply(to=command1)
|
||||
tplaybook >> tutils.reply(to=command2)
|
||||
else:
|
||||
tplaybook >> tutils.reply(to=command2)
|
||||
tplaybook >> tutils.reply(to=command1)
|
||||
assert tplaybook
|
||||
assert a() == 1
|
||||
assert b() == 2
|
||||
|
||||
@ -1,5 +1,4 @@
|
||||
import collections.abc
|
||||
import copy
|
||||
import difflib
|
||||
import itertools
|
||||
import re
|
||||
@ -158,7 +157,18 @@ class Playbook:
|
||||
if c is None:
|
||||
return self
|
||||
assert isinstance(c, commands.Command)
|
||||
self.expected.append(c)
|
||||
|
||||
prev = self.expected[-1]
|
||||
two_subsequent_sends_to_the_same_remote = (
|
||||
isinstance(c, commands.SendData)
|
||||
and isinstance(prev, commands.SendData)
|
||||
and prev.connection is c.connection
|
||||
)
|
||||
if two_subsequent_sends_to_the_same_remote:
|
||||
prev.data += c.data
|
||||
else:
|
||||
self.expected.append(c)
|
||||
|
||||
return self
|
||||
|
||||
def __bool__(self):
|
||||
@ -191,17 +201,18 @@ class Playbook:
|
||||
self.actual.append(_TracebackInPlaybook(traceback.format_exc()))
|
||||
break
|
||||
self.actual.extend(cmds)
|
||||
pos = i
|
||||
for cmd in cmds:
|
||||
pos += 1
|
||||
assert self.actual[pos] == cmd
|
||||
if isinstance(cmd, commands.CloseConnection):
|
||||
if cmd.half_close:
|
||||
cmd.connection.state &= ~ConnectionState.CAN_WRITE
|
||||
else:
|
||||
cmd.connection.state = ConnectionState.CLOSED
|
||||
if not self.logs:
|
||||
for offset, cmd in enumerate(cmds):
|
||||
pos = i + 1 + offset
|
||||
elif isinstance(cmd, commands.Log):
|
||||
need_to_emulate_log = (
|
||||
isinstance(cmd, commands.Log) and
|
||||
not self.logs and
|
||||
cmd.level in ("debug", "info") and
|
||||
(
|
||||
pos >= len(self.expected)
|
||||
@ -210,25 +221,34 @@ class Playbook:
|
||||
)
|
||||
if need_to_emulate_log:
|
||||
self.expected.insert(pos, cmd)
|
||||
if not self.hooks:
|
||||
last_cmd = self.actual[-1]
|
||||
pos = i + len(cmds)
|
||||
need_to_emulate_hook = (
|
||||
isinstance(last_cmd, commands.Hook) and
|
||||
(
|
||||
pos >= len(self.expected) or
|
||||
(not (isinstance(self.expected[pos], commands.Hook)
|
||||
and self.expected[pos].name == last_cmd.name))
|
||||
)
|
||||
)
|
||||
if need_to_emulate_hook:
|
||||
self.expected.insert(pos, last_cmd)
|
||||
self.expected.insert(pos + 1, events.HookReply(last_cmd))
|
||||
elif isinstance(cmd, commands.Hook) and not self.hooks:
|
||||
need_to_emulate_hook = (
|
||||
not self.hooks
|
||||
and (
|
||||
pos >= len(self.expected) or
|
||||
(not (
|
||||
isinstance(self.expected[pos], commands.Hook)
|
||||
and self.expected[pos].name == cmd.name
|
||||
))
|
||||
)
|
||||
)
|
||||
if need_to_emulate_hook:
|
||||
self.expected.insert(pos, cmd)
|
||||
if cmd.blocking:
|
||||
self.expected.insert(pos + 1, events.HookReply(cmd))
|
||||
elif isinstance(cmd, commands.SendData):
|
||||
prev = self.actual[pos - 1]
|
||||
two_subsequent_sends_to_the_same_remote = (
|
||||
isinstance(prev, commands.SendData) and
|
||||
cmd.connection is prev.connection
|
||||
)
|
||||
if two_subsequent_sends_to_the_same_remote:
|
||||
prev.data += cmd.data
|
||||
self.actual.pop(pos)
|
||||
pos -= 1
|
||||
eq(self.expected[i:], self.actual[i:]) # compare now already to set placeholders
|
||||
i += 1
|
||||
|
||||
self.actual = _merge_sends(self.actual)
|
||||
self.expected = _merge_sends(self.expected)
|
||||
|
||||
if not eq(self.expected, self.actual):
|
||||
self._errored = True
|
||||
diff = "\n".join(difflib.ndiff(
|
||||
@ -247,13 +267,6 @@ class Playbook:
|
||||
if is_final_destruct or (not self._errored and len(self.actual) < len(self.expected)):
|
||||
raise RuntimeError("Unfinished playbook!")
|
||||
|
||||
def fork(self):
|
||||
"""
|
||||
Fork the current playbook to assert a second execution stream from here on.
|
||||
Returns a new playbook instance.
|
||||
"""
|
||||
return copy.deepcopy(self)
|
||||
|
||||
|
||||
class reply(events.Event):
|
||||
args: typing.Tuple[typing.Any, ...]
|
||||
|
||||
Loading…
Reference in New Issue
Block a user