Skip to content

Commit

Permalink
Merge pull request #336 from tohojo/optional-headers
Browse files Browse the repository at this point in the history
HeaderMatchingFilter: Gracefully handle missing headers
  • Loading branch information
GuillaumeSeren authored Nov 21, 2023
2 parents e1f6488 + 38797f6 commit 3e55b17
Show file tree
Hide file tree
Showing 2 changed files with 96 additions and 7 deletions.
19 changes: 12 additions & 7 deletions afew/filters/HeaderMatchingFilter.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@

from afew.filters.BaseFilter import Filter

from notmuch.errors import NullPointerError

import re


Expand All @@ -22,10 +24,13 @@ def __init__(self, database, **kwargs):
def handle_message(self, message):
if self.header is not None and self.pattern is not None:
if not self._tag_blacklist.intersection(message.get_tags()):
value = message.get_header(self.header)
match = self.pattern.search(value)
if match:
tagdict = {k: v.lower() for k, v in match.groupdict().items()}
sub = (lambda tag: tag.format(**tagdict))
self.remove_tags(message, *map(sub, self._tags_to_remove))
self.add_tags(message, *map(sub, self._tags_to_add))
try:
value = message.get_header(self.header)
match = self.pattern.search(value)
if match:
tagdict = {k: v.lower() for k, v in match.groupdict().items()}
sub = (lambda tag: tag.format(**tagdict))
self.remove_tags(message, *map(sub, self._tags_to_remove))
self.add_tags(message, *map(sub, self._tags_to_add))
except NullPointerError:
pass
84 changes: 84 additions & 0 deletions afew/tests/test_headermatchingfilter.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
"""Test suite for DKIMValidityFilter.
"""
import unittest
from email.utils import make_msgid
from unittest import mock

from afew.Database import Database
from afew.filters.HeaderMatchingFilter import HeaderMatchingFilter

from notmuch.errors import NullPointerError


class _AddTags: # pylint: disable=too-few-public-methods
"""Mock for `add_tags` method of base filter. We need to easily collect
tags added by filter for test assertion.
"""
def __init__(self, tags):
self._tags = tags

def __call__(self, message, *tags):
self._tags.update(tags)


def _make_header_matching_filter():
"""Make `HeaderMatchingFilter` with mocked `HeaderMatchingFilter.add_tags`
method, so in tests we can easily check what tags were added by filter
without fiddling with db.
"""
tags = set()
add_tags = _AddTags(tags)
header_filter = HeaderMatchingFilter(Database(), header="X-test", pattern="")
header_filter.add_tags = add_tags
return header_filter, tags


def _make_message(should_fail):
"""Make mock email Message.
Mocked methods:
- `get_header()` returns non-empty string. When testing with mocked
function for verifying DKIM signature, DKIM signature doesn't matter as
long as it's non-empty string.
- `get_filenames()` returns list of non-empty string. When testing with
mocked file open, it must just be non-empty string.
- `get_message_id()` returns some generated message ID.
"""
message = mock.Mock()
if should_fail:
message.get_header.side_effect = NullPointerError
else:
message.get_header.return_value = 'header'
message.get_filenames.return_value = ['a']
message.get_tags.return_value = ['a']
message.get_message_id.return_value = make_msgid()
return message


class TestHeaderMatchingFilter(unittest.TestCase):
"""Test suite for `HeaderMatchingFilter`.
"""
@mock.patch('afew.filters.HeaderMatchingFilter.open',
mock.mock_open(read_data=b''))
def test_header_exists(self):
"""Test message with header that exists.
"""
header_filter, tags = _make_header_matching_filter()
message = _make_message(False)
header_filter.handle_message(message)

self.assertSetEqual(tags, set())

@mock.patch('afew.filters.HeaderMatchingFilter.open',
mock.mock_open(read_data=b''))
def test_header_doesnt_exist(self):
"""Test message with header that exists.
"""
header_filter, tags = _make_header_matching_filter()
message = _make_message(True)
header_filter.handle_message(message)

self.assertSetEqual(tags, set())

0 comments on commit 3e55b17

Please sign in to comment.