Skip to content

Commit d75af1d

Browse files
xmunozwoodruffw
authored andcommitted
Implement scheduled checks #7093 (#7271)
* Implement scheduled checks #7093 - Rename `run_backfill` to `run_evaluation` in admin malware view - Modify `run` and `scan` method signatures to accept `**kwargs` - Extend `run_check` to accomodate scheduled check functionality * Reduce unit test flakiness * Code review changes. Also replace `check.hooked_object` with `check.hooked_object.value` in check detail template. * tests, warehouse: enum fixes * Fix lint error Co-authored-by: William Woodruff <[email protected]>
1 parent f2b93df commit d75af1d

File tree

12 files changed

+197
-67
lines changed

12 files changed

+197
-67
lines changed

tests/common/checks/scheduled.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -27,11 +27,11 @@ class ExampleScheduledCheck(MalwareCheckBase):
2727
def __init__(self, db):
2828
super().__init__(db)
2929

30-
def scan(self):
30+
def scan(self, **kwargs):
3131
project = self.db.query(Project).first()
3232
self.add_verdict(
3333
project_id=project.id,
34-
classification=VerdictClassification.benign,
34+
classification=VerdictClassification.Benign,
3535
confidence=VerdictConfidence.High,
3636
message="Nothing to see here!",
3737
)

tests/unit/admin/test_routes.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -133,8 +133,8 @@ def test_includeme():
133133
domain=warehouse,
134134
),
135135
pretend.call(
136-
"admin.checks.run_backfill",
137-
"/admin/checks/{check_name}/run_backfill",
136+
"admin.checks.run_evaluation",
137+
"/admin/checks/{check_name}/run_evaluation",
138138
domain=warehouse,
139139
),
140140
pretend.call("admin.verdicts.list", "/admin/verdicts/", domain=warehouse),

tests/unit/admin/views/test_checks.py

+40-19
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,8 @@
1616
from pyramid.httpexceptions import HTTPNotFound
1717

1818
from warehouse.admin.views import checks as views
19-
from warehouse.malware.models import MalwareCheckState
19+
from warehouse.malware.models import MalwareCheckState, MalwareCheckType
20+
from warehouse.malware.tasks import backfill, run_check
2021

2122
from ....common.db.malware import MalwareCheckFactory
2223

@@ -46,6 +47,7 @@ def test_get_check(self, db_request):
4647
"check": check,
4748
"checks": [check],
4849
"states": MalwareCheckState,
50+
"evaluation_run_size": 10000,
4951
}
5052

5153
def test_get_check_many_versions(self, db_request):
@@ -56,6 +58,7 @@ def test_get_check_many_versions(self, db_request):
5658
"check": check2,
5759
"checks": [check2, check1],
5860
"states": MalwareCheckState,
61+
"evaluation_run_size": 10000,
5962
}
6063

6164
def test_get_check_not_found(self, db_request):
@@ -129,17 +132,17 @@ def test_change_to_invalid_state(self, db_request):
129132
assert check.state == initial_state
130133

131134

132-
class TestRunBackfill:
135+
class TestRunEvaluation:
133136
@pytest.mark.parametrize(
134137
("check_state", "message"),
135138
[
136139
(
137140
MalwareCheckState.Disabled,
138-
"Check must be in 'enabled' or 'evaluation' state to run a backfill.",
141+
"Check must be in 'enabled' or 'evaluation' state to manually execute.",
139142
),
140143
(
141144
MalwareCheckState.WipedOut,
142-
"Check must be in 'enabled' or 'evaluation' state to run a backfill.",
145+
"Check must be in 'enabled' or 'evaluation' state to manually execute.",
143146
),
144147
],
145148
)
@@ -152,23 +155,29 @@ def test_invalid_backfill_parameters(self, db_request, check_state, message):
152155
)
153156

154157
db_request.route_path = pretend.call_recorder(
155-
lambda *a, **kw: "/admin/checks/%s/run_backfill" % check.name
158+
lambda *a, **kw: "/admin/checks/%s/run_evaluation" % check.name
156159
)
157160

158-
views.run_backfill(db_request)
161+
views.run_evaluation(db_request)
159162

160163
assert db_request.session.flash.calls == [pretend.call(message, queue="error")]
161164

162-
def test_sucess(self, db_request):
163-
check = MalwareCheckFactory.create(state=MalwareCheckState.Enabled)
165+
@pytest.mark.parametrize(
166+
("check_type"), [MalwareCheckType.EventHook, MalwareCheckType.Scheduled]
167+
)
168+
def test_success(self, db_request, check_type):
169+
170+
check = MalwareCheckFactory.create(
171+
check_type=check_type, state=MalwareCheckState.Enabled
172+
)
164173
db_request.matchdict["check_name"] = check.name
165174

166175
db_request.session = pretend.stub(
167176
flash=pretend.call_recorder(lambda *a, **kw: None)
168177
)
169178

170179
db_request.route_path = pretend.call_recorder(
171-
lambda *a, **kw: "/admin/checks/%s/run_backfill" % check.name
180+
lambda *a, **kw: "/admin/checks/%s/run_evaluation" % check.name
172181
)
173182

174183
backfill_recorder = pretend.stub(
@@ -177,13 +186,25 @@ def test_sucess(self, db_request):
177186

178187
db_request.task = pretend.call_recorder(lambda *a, **kw: backfill_recorder)
179188

180-
views.run_backfill(db_request)
181-
182-
assert db_request.session.flash.calls == [
183-
pretend.call(
184-
"Running %s on 10000 %ss!" % (check.name, check.hooked_object.value),
185-
queue="success",
186-
)
187-
]
188-
189-
assert backfill_recorder.delay.calls == [pretend.call(check.name, 10000)]
189+
views.run_evaluation(db_request)
190+
191+
if check_type == MalwareCheckType.EventHook:
192+
assert db_request.session.flash.calls == [
193+
pretend.call(
194+
"Running %s on 10000 %ss!"
195+
% (check.name, check.hooked_object.value),
196+
queue="success",
197+
)
198+
]
199+
assert db_request.task.calls == [pretend.call(backfill)]
200+
assert backfill_recorder.delay.calls == [pretend.call(check.name, 10000)]
201+
elif check_type == MalwareCheckType.Scheduled:
202+
assert db_request.session.flash.calls == [
203+
pretend.call("Running %s now!" % check.name, queue="success",)
204+
]
205+
assert db_request.task.calls == [pretend.call(run_check)]
206+
assert backfill_recorder.delay.calls == [
207+
pretend.call(check.name, manually_triggered=True)
208+
]
209+
else:
210+
raise Exception("Invalid check type: %s" % check_type)

tests/unit/malware/test_init.py

+10
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,12 @@
1414

1515
import pretend
1616

17+
from celery.schedules import crontab
18+
1719
from warehouse import malware
1820
from warehouse.malware import utils
1921
from warehouse.malware.interfaces import IMalwareCheckService
22+
from warehouse.malware.tasks import run_check
2023

2124
from ...common import checks as test_checks
2225
from ...common.db.accounts import UserFactory
@@ -165,10 +168,17 @@ def test_includeme(monkeypatch):
165168
registry=pretend.stub(
166169
settings={"malware_check.backend": "TestMalwareCheckService"}
167170
),
171+
add_periodic_task=pretend.call_recorder(lambda *a, **kw: None),
168172
)
169173

170174
malware.includeme(config)
171175

172176
assert config.register_service_factory.calls == [
173177
pretend.call(malware_check_class.create_service, IMalwareCheckService)
174178
]
179+
180+
assert config.add_periodic_task.calls == [
181+
pretend.call(
182+
crontab(minute="0", hour="*/8"), run_check, args=("ExampleScheduledCheck",)
183+
)
184+
]

tests/unit/malware/test_tasks.py

+54-17
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,6 @@
1414
import pretend
1515
import pytest
1616

17-
from sqlalchemy.orm.exc import NoResultFound
18-
1917
from warehouse.malware import tasks
2018
from warehouse.malware.models import MalwareCheck, MalwareCheckState, MalwareVerdict
2119

@@ -34,45 +32,86 @@ def test_success(self, db_request, monkeypatch):
3432
name="ExampleHookedCheck", state=MalwareCheckState.Enabled
3533
)
3634
task = pretend.stub()
37-
tasks.run_check(task, db_request, "ExampleHookedCheck", file0.id)
35+
tasks.run_check(task, db_request, "ExampleHookedCheck", obj_id=file0.id)
3836

3937
assert db_request.route_url.calls == [
4038
pretend.call("packaging.file", path=file0.path)
4139
]
4240
assert db_request.db.query(MalwareVerdict).one()
4341

44-
def test_disabled_check(self, db_request, monkeypatch):
42+
@pytest.mark.parametrize(("manually_triggered"), [True, False])
43+
def test_evaluation_run(self, db_session, monkeypatch, manually_triggered):
44+
monkeypatch.setattr(tasks, "checks", test_checks)
45+
MalwareCheckFactory.create(
46+
name="ExampleScheduledCheck", state=MalwareCheckState.Evaluation
47+
)
48+
ProjectFactory.create()
49+
task = pretend.stub()
50+
51+
request = pretend.stub(
52+
db=db_session,
53+
log=pretend.stub(info=pretend.call_recorder(lambda *args, **kwargs: None)),
54+
)
55+
56+
tasks.run_check(
57+
task,
58+
request,
59+
"ExampleScheduledCheck",
60+
manually_triggered=manually_triggered,
61+
)
62+
63+
if manually_triggered:
64+
assert db_session.query(MalwareVerdict).one()
65+
else:
66+
assert request.log.info.calls == [
67+
pretend.call(
68+
"ExampleScheduledCheck is in the `evaluation` state and must be \
69+
manually triggered to run."
70+
)
71+
]
72+
assert db_session.query(MalwareVerdict).all() == []
73+
74+
def test_disabled_check(self, db_session, monkeypatch):
4575
monkeypatch.setattr(tasks, "checks", test_checks)
4676
MalwareCheckFactory.create(
4777
name="ExampleHookedCheck", state=MalwareCheckState.Disabled
4878
)
4979
task = pretend.stub()
80+
request = pretend.stub(
81+
db=db_session,
82+
log=pretend.stub(info=pretend.call_recorder(lambda *args, **kwargs: None)),
83+
)
5084

5185
file = FileFactory.create()
5286

53-
with pytest.raises(NoResultFound):
54-
tasks.run_check(task, db_request, "ExampleHookedCheck", file.id)
87+
tasks.run_check(
88+
task, request, "ExampleHookedCheck", obj_id=file.id,
89+
)
90+
91+
assert request.log.info.calls == [
92+
pretend.call("Check ExampleHookedCheck isn't active. Aborting.")
93+
]
5594

5695
def test_missing_check(self, db_request, monkeypatch):
5796
monkeypatch.setattr(tasks, "checks", test_checks)
5897
task = pretend.stub()
5998

60-
file = FileFactory.create()
61-
6299
with pytest.raises(AttributeError):
63-
tasks.run_check(task, db_request, "DoesNotExistCheck", file.id)
100+
tasks.run_check(
101+
task, db_request, "DoesNotExistCheck",
102+
)
64103

65104
def test_retry(self, db_session, monkeypatch):
105+
monkeypatch.setattr(tasks, "checks", test_checks)
66106
exc = Exception("Scan failed")
67107

68108
def scan(self, **kwargs):
69109
raise exc
70110

71-
monkeypatch.setattr(tasks, "checks", test_checks)
72111
monkeypatch.setattr(tasks.checks.ExampleHookedCheck, "scan", scan)
73112

74113
MalwareCheckFactory.create(
75-
name="ExampleHookedCheck", state=MalwareCheckState.Evaluation
114+
name="ExampleHookedCheck", state=MalwareCheckState.Enabled
76115
)
77116

78117
task = pretend.stub(
@@ -87,7 +126,7 @@ def scan(self, **kwargs):
87126
file = FileFactory.create()
88127

89128
with pytest.raises(celery.exceptions.Retry):
90-
tasks.run_check(task, request, "ExampleHookedCheck", file.id)
129+
tasks.run_check(task, request, "ExampleHookedCheck", obj_id=file.id)
91130

92131
assert request.log.error.calls == [
93132
pretend.call("Error executing check ExampleHookedCheck: Scan failed")
@@ -108,9 +147,8 @@ def test_invalid_check_name(self, db_request, monkeypatch):
108147
)
109148
def test_run(self, db_session, num_objects, num_runs, monkeypatch):
110149
monkeypatch.setattr(tasks, "checks", test_checks)
111-
files = []
112150
for i in range(num_objects):
113-
files.append(FileFactory.create())
151+
FileFactory.create()
114152

115153
MalwareCheckFactory.create(
116154
name="ExampleHookedCheck", state=MalwareCheckState.Enabled
@@ -133,15 +171,14 @@ def test_run(self, db_session, num_objects, num_runs, monkeypatch):
133171
pretend.call("Running backfill on %d Files." % num_runs)
134172
]
135173

136-
assert enqueue_recorder.delay.calls == [
137-
pretend.call("ExampleHookedCheck", files[i].id) for i in range(num_runs)
138-
]
174+
assert len(enqueue_recorder.delay.calls) == num_runs
139175

140176

141177
class TestSyncChecks:
142178
def test_no_updates(self, db_session, monkeypatch):
143179
monkeypatch.setattr(tasks, "checks", test_checks)
144180
monkeypatch.setattr(tasks.checks.ExampleScheduledCheck, "version", 2)
181+
145182
MalwareCheckFactory.create(
146183
name="ExampleHookedCheck", state=MalwareCheckState.Disabled
147184
)

warehouse/admin/routes.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -140,8 +140,8 @@ def includeme(config):
140140
domain=warehouse,
141141
)
142142
config.add_route(
143-
"admin.checks.run_backfill",
144-
"/admin/checks/{check_name}/run_backfill",
143+
"admin.checks.run_evaluation",
144+
"/admin/checks/{check_name}/run_evaluation",
145145
domain=warehouse,
146146
)
147147
config.add_route("admin.verdicts.list", "/admin/verdicts/", domain=warehouse)

warehouse/admin/templates/admin/malware/checks/detail.html

+16-2
Original file line numberDiff line numberDiff line change
@@ -30,12 +30,22 @@ <h4>Revision History</h4>
3030
<tr>
3131
<th>Version</th>
3232
<th>State</th>
33+
{% if check.check_type.value == "event_hook" %}
34+
<th>Hooked Object</th>
35+
{% else %}
36+
<th>Schedule</th>
37+
{% endif %}
3338
<th>Created</th>
3439
</tr>
3540
{% for c in checks %}
3641
<tr>
3742
<td>{{ c.version }}</td>
3843
<td>{{ c.state.value }}</td>
44+
{% if check.check_type.value == "event_hook" %}
45+
<td>{{ c.hooked_object.value }}</td>
46+
{% else %}
47+
<td><pre>{{ c.schedule }}</pre></td>
48+
{% endif %}
3949
<td>{{ c.created }}</td>
4050
</tr>
4151
{% endfor %}
@@ -69,10 +79,14 @@ <h3 class="box-title">Change State</h3>
6979
<div class="box-header with-border">
7080
<h3 class="box-title">Run Evaluation</h3>
7181
</div>
72-
<form method="POST" action="{{ request.route_path('admin.checks.run_backfill', check_name=check.name) }}">
82+
<form method="POST" action="{{ request.route_path('admin.checks.run_evaluation', check_name=check.name) }}">
7383
<input name="csrf_token" type="hidden" value="{{ request.session.get_csrf_token() }}">
7484
<div class="box-body">
75-
<p>Run this check against 10,000 {{ check.hooked_object.value }}s, selected at random. This is used to evaluate the efficacy of a check.</p>
85+
{% if check.check_type.value == "event_hook" %}
86+
<p>Run this check against {{ evaluation_run_size }} {{ check.hooked_object.value }}s, selected at random. This is used to evaluate the efficacy of a check.</p>
87+
{% else %}
88+
<p>Execute this check now.</p>
89+
{% endif %}
7690
<div class="pull-right col-sm-4">
7791
<button type="submit" class="btn btn-primary pull-right">Run</button>
7892
</div>

warehouse/admin/templates/admin/malware/checks/index.html

+2-2
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
<tr>
2727
<th>Check Name</th>
2828
<th>State</th>
29-
<th>Revisions</th>
29+
<th>Type</th>
3030
<th>Last Modified</th>
3131
<th>Description</th>
3232
</tr>
@@ -38,7 +38,7 @@
3838
</a>
3939
</td>
4040
<td>{{ check.state.value }}</td>
41-
<td>{{ check.version }}</td>
41+
<td>{{ check.check_type.value }}</td>
4242
<td>{{ check.created }}</td>
4343
<td>{{ check.short_description }}</td>
4444
</tr>

0 commit comments

Comments
 (0)