/
compute.py
executable file
·82 lines (68 loc) · 2.48 KB
/
compute.py
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
81
82
import os
import requests_mock
import json
import uuid
from rest_framework.test import APIClient
from webapp.apps.comp.compute import Compute
class CallbackException(Exception):
pass
class MockCompute(Compute):
outputs = None
client = None
sim = None
user = "comp-api-user"
password = "heyhey2222"
def __init__(self, num_times_to_wait=0):
self.count = 0
self.num_times_to_wait = num_times_to_wait
def remote_submit_job(self, url, data, timeout, headers=None):
print("mocking:", url)
with requests_mock.Mocker() as mock:
resp = {"job_id": str(uuid.uuid4()), "qlength": 2}
resp = json.dumps(resp)
mock.register_uri("POST", url, text=resp)
self.last_posted = data
return Compute.remote_submit_job(self, url, data, timeout)
def remote_query_job(self, url):
# Need to login as the comp-api-user
self.client.login(username=self.user, password=self.password)
if isinstance(self.client, APIClient):
format_kwarg = {"format": "json"}
else:
format_kwarg = {"content_type": "application/json"}
resp = self.client.put(
"/outputs/api/",
data=dict(json.loads(self.outputs), **{"job_id": self.sim.job_id}),
**format_kwarg,
)
if resp.status_code != 200:
raise CallbackException(
f"Status code: {resp.status_code}\n {json.dumps(resp.data, indent=4)}"
)
self.client = None
self.sim = None
with requests_mock.Mocker() as mock:
text = "NO"
mock.register_uri("GET", url, text=text)
return Compute.remote_query_job(self, url)
def remote_get_job(self, url):
self.count += 1
with requests_mock.Mocker() as mock:
mock.register_uri("GET", url, text=self.outputs)
return Compute.remote_get_job(self, url)
def reset_count(self):
"""
reset worker node count
"""
self.count = 0
class MockComputeWorkerFailure(MockCompute):
next_response = None
outputs = json.dumps({"status": "WORKER_FAILURE", "traceback": "Error: whoops"})
def remote_query_job(self, url):
self.client = None
self.sim = None
with requests_mock.Mocker() as mock:
print("mocking: ", url)
text = "FAIL"
mock.register_uri("GET", url, text=text)
return Compute.remote_query_job(self, url)