tracker/clients/python/tests/test_get_dmarc.py at flux-sync · LearningAzure2020/tracker · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
"""Tests for functions that get DMARC summaries for a domain"""
import json

from gql import Client

# Avoid confusion with gql Client
import tracker_client.dmarc as tc
import tracker_client.queries as queries

# pylint: disable=no-member


def test_get_dmarc_summary_error(mocker, error_message):
    """Test that get_dmarc_summary correctly handles error response"""
    mocker.patch(
        "tracker_client.dmarc.execute_query",
        auto_spec=True,
        return_value=error_message,
    )

    client = Client()
    result = tc.get_dmarc_summary(client, "foo.bar", "september", "2020")

    assert result == json.dumps(
        error_message,
        indent=4,
    )


def test_get_dmarc_summary(mocker, monthly_dmarc_input, monthly_dmarc_output):
    """Test that get_dmarc_summary produces correct output"""
    mocker.patch(
        "tracker_client.dmarc.execute_query",
        auto_spec=True,
        return_value=monthly_dmarc_input,
    )

    expected_params = {"domain": "foo.bar", "month": "SEPTEMBER", "year": "2020"}
    client = Client()
    result = tc.get_dmarc_summary(client, "foo.bar", "september", "2020")

    tc.execute_query.assert_called_once_with(
        client, queries.DMARC_SUMMARY, expected_params
    )
    assert result == json.dumps(monthly_dmarc_output, indent=4)


def test_get_yearly_dmarc_summaries_error(mocker, error_message):
    """Test that get_yearly_dmarc_summaries correctly handles error response"""
    mocker.patch(
        "tracker_client.dmarc.execute_query",
        auto_spec=True,
        return_value=error_message,
    )

    client = Client()
    result = tc.get_yearly_dmarc_summaries(client, "foo.bar")

    assert result == json.dumps(
        error_message,
        indent=4,
    )


def test_get_yearly_dmarc_summaries(mocker, yearly_dmarc_input, yearly_dmarc_output):
    """Test that get_yearly_dmarc_summaries produces correct output"""
    mocker.patch(
        "tracker_client.dmarc.execute_query",
        auto_spec=True,
        return_value=yearly_dmarc_input,
    )

    expected_params = {"domain": "foo.bar"}
    client = Client()
    result = tc.get_yearly_dmarc_summaries(client, "foo.bar")

    tc.execute_query.assert_called_once_with(
        client, queries.DMARC_YEARLY_SUMMARIES, expected_params
    )
    assert result == json.dumps(yearly_dmarc_output, indent=4)