Skip to content

Commit ba9a4eb

Browse files
authored
added timings, tests (#75)
Signed-off-by: dan nelson <[email protected]>
1 parent 6a7739b commit ba9a4eb

File tree

2 files changed

+71
-0
lines changed

2 files changed

+71
-0
lines changed

replicate/prediction.py

+3
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,9 @@ class Prediction(BaseModel):
1717
output: Optional[Any]
1818
status: str
1919
version: Optional[Version]
20+
started_at: Optional[str]
21+
created_at: Optional[str]
22+
completed_at: Optional[str]
2023

2124
def wait(self):
2225
"""Wait for prediction to finish."""

tests/test_prediction.py

+68
Original file line numberDiff line numberDiff line change
@@ -48,3 +48,71 @@ def test_cancel():
4848
rsp = responses.post("https://api.replicate.com/v1/predictions/p1/cancel", json={})
4949
prediction.cancel()
5050
assert rsp.call_count == 1
51+
52+
53+
@responses.activate
54+
def test_async_timings():
55+
client = create_client()
56+
version = create_version(client)
57+
58+
responses.post(
59+
"https://api.replicate.com/v1/predictions",
60+
match=[
61+
matchers.json_params_matcher(
62+
{
63+
"version": "v1",
64+
"input": {"text": "hello"},
65+
"webhook_completed": "https://example.com/webhook",
66+
}
67+
),
68+
],
69+
json={
70+
"id": "p1",
71+
"version": "v1",
72+
"urls": {
73+
"get": "https://api.replicate.com/v1/predictions/p1",
74+
"cancel": "https://api.replicate.com/v1/predictions/p1/cancel",
75+
},
76+
"created_at": "2022-04-26T20:00:40.658234Z",
77+
"source": "api",
78+
"status": "processing",
79+
"input": {"text": "hello"},
80+
"output": None,
81+
"error": None,
82+
"logs": "",
83+
},
84+
)
85+
86+
responses.get(
87+
"https://api.replicate.com/v1/predictions/p1",
88+
json={
89+
"id": "p1",
90+
"version": "v1",
91+
"urls": {
92+
"get": "https://api.replicate.com/v1/predictions/p1",
93+
"cancel": "https://api.replicate.com/v1/predictions/p1/cancel",
94+
},
95+
"created_at": "2022-04-26T20:00:40.658234Z",
96+
"completed_at": "2022-04-26T20:02:27.648305Z",
97+
"source": "api",
98+
"status": "succeeded",
99+
"input": {"text": "hello"},
100+
"output": "hello world",
101+
"error": None,
102+
"logs": "",
103+
},
104+
)
105+
106+
prediction = client.predictions.create(
107+
version=version,
108+
input={"text": "hello"},
109+
webhook_completed="https://example.com/webhook",
110+
)
111+
112+
assert prediction.created_at == "2022-04-26T20:00:40.658234Z"
113+
assert prediction.completed_at == None
114+
assert prediction.output == None
115+
prediction.wait()
116+
assert prediction.created_at == "2022-04-26T20:00:40.658234Z"
117+
assert prediction.completed_at == "2022-04-26T20:02:27.648305Z"
118+
assert prediction.output == "hello world"

0 commit comments

Comments
 (0)