1import sys
2
3import pytest
4
5import postfix_mta_sts_resolver.__main__ as main
6
7class MockCmdline:
8    def __init__(self, *args):
9        self._cmdline = args
10
11    def __enter__(self):
12        self._old_cmdline = sys.argv
13        sys.argv = list(self._cmdline)
14
15    def __exit__(self, exc_type, exc_value, traceback):
16        sys.argv = self._old_cmdline
17
18def test_parse_args():
19    with MockCmdline("mta-sts-query", "example.com"):
20        args = main.parse_args()
21    assert args.domain == 'example.com'
22    assert args.known_version is None
23
24def test_parse_args_with_version():
25    with MockCmdline("mta-sts-query", "example.com", "123"):
26        args = main.parse_args()
27    assert args.domain == 'example.com'
28    assert args.known_version == "123"
29