Skip to content

Commit

Permalink
docs: update test_lit_server.py (#222)
Browse files Browse the repository at this point in the history
Indentity -> Identity
  • Loading branch information
eltociear authored Aug 26, 2024
1 parent 7923bc4 commit 0a6045a
Showing 1 changed file with 6 additions and 6 deletions.
12 changes: 6 additions & 6 deletions tests/test_lit_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -356,7 +356,7 @@ def test_server_terminate():
mock_manager.shutdown.assert_called()


class IndentityAPI(ls.examples.SimpleLitAPI):
class IdentityAPI(ls.examples.SimpleLitAPI):
def predict(self, x, context):
context["input"] = x
return self.model(x)
Expand All @@ -366,7 +366,7 @@ def encode_response(self, output, context):
return {"output": input}


class IndentityBatchedAPI(ls.examples.SimpleBatchedAPI):
class IdentityBatchedAPI(ls.examples.SimpleBatchedAPI):
def predict(self, x_batch, context):
for c, x in zip(context, x_batch):
c["input"] = x
Expand All @@ -377,7 +377,7 @@ def encode_response(self, output, context):
return {"output": input}


class IndentityBatchedStreamingAPI(ls.examples.SimpleBatchedAPI):
class IdentityBatchedStreamingAPI(ls.examples.SimpleBatchedAPI):
def predict(self, x_batch, context):
for c, x in zip(context, x_batch):
c["input"] = x
Expand Down Expand Up @@ -407,22 +407,22 @@ def dummy_load_and_raise(resp):
mocked_load_and_raise.side_effect = dummy_load_and_raise

# Test context injection with single loop
api = IndentityAPI()
api = IdentityAPI()
server = LitServer(api)
with wrap_litserve_start(server) as server:
async with LifespanManager(server.app) as manager, AsyncClient(app=manager.app, base_url="http://test") as ac:
resp = await ac.post("/predict", json={"input": 5.0}, timeout=10)
assert resp.json()["output"] == 5.0, "output from Identity server must be same as input"

# Test context injection with batched loop
server = LitServer(IndentityBatchedAPI(), max_batch_size=2, batch_timeout=0.01)
server = LitServer(IdentityBatchedAPI(), max_batch_size=2, batch_timeout=0.01)
with wrap_litserve_start(server) as server:
async with LifespanManager(server.app) as manager, AsyncClient(app=manager.app, base_url="http://test") as ac:
resp = await ac.post("/predict", json={"input": 5.0}, timeout=10)
assert resp.json()["output"] == 5.0, "output from Identity server must be same as input"

# Test context injection with batched streaming loop
server = LitServer(IndentityBatchedStreamingAPI(), max_batch_size=2, batch_timeout=0.01, stream=True)
server = LitServer(IdentityBatchedStreamingAPI(), max_batch_size=2, batch_timeout=0.01, stream=True)
with wrap_litserve_start(server) as server:
async with LifespanManager(server.app) as manager, AsyncClient(app=manager.app, base_url="http://test") as ac:
resp = await ac.post("/predict", json={"input": 5.0}, timeout=10)
Expand Down

0 comments on commit 0a6045a

Please sign in to comment.