forked from mozilla-releng/balrog
-
Notifications
You must be signed in to change notification settings - Fork 5
/
fakes.py
67 lines (48 loc) · 1.85 KB
/
fakes.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
import json
from auslib.db import GCSHistory, GCSHistoryAsync
def FakeBlobFactory(exc=None):
class fb:
def __init__(self, data=None):
self.data = data
async def upload(self, body, *args, **kwargs):
if exc:
raise exc("I failed to upload")
else:
self.data = body
def upload_from_string(self, body, *args, **kwargs):
if exc:
raise exc("I failed to upload")
else:
self.data = body
return fb
FakeBlob = FakeBlobFactory()
class FakeBucket:
def __init__(self):
self.blobs = {}
def new_blob(self, blob_name, *args, **kwargs):
self.blobs[blob_name] = FakeBlob()
return self.blobs[blob_name]
blob = new_blob
class FakeGCSHistory(GCSHistory):
def __init__(self, *args, identifier_columns=["name"], **kwargs):
self.bucket = FakeBucket()
self.identifier_columns = identifier_columns
self.data_column = "data"
def _getBucket(self, identifier):
return lambda use_gcloud_aio: self.bucket
def getChange(self, change_id=None, column_values=None, data_version=None, transaction=None):
name = column_values["name"]
prefix = f"{name}/{data_version}-"
data = None
for bname, blob in self.bucket.blobs.items():
if bname.startswith(prefix):
data = blob.data
break
return {"name": column_values["name"], "data_version": data_version, "data": json.loads(data)}
class FakeGCSHistoryAsync(GCSHistoryAsync):
def __init__(self, *args, identifier_columns=["name"], **kwargs):
self.bucket = FakeBucket()
self.identifier_columns = identifier_columns
self.data_column = "data"
def _getBucket(self, identifier):
return lambda session: self.bucket