Skip to content

Commit

Permalink
fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
d80tb7 committed Sep 25, 2024
1 parent d284e1a commit 6f38664
Showing 1 changed file with 7 additions and 6 deletions.
13 changes: 7 additions & 6 deletions client/python/armada_client/asyncio_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,7 @@ def __init__(
event_timeout: timedelta = timedelta(minutes=15),
) -> None:
self.submit_stub = submit_pb2_grpc.SubmitStub(channel)
self.queue_stub = submit_pb2_grpc.QueueServiceStub(channel)
self.event_stub = event_pb2_grpc.EventStub(channel)
self.job_stub = job_pb2_grpc.JobsStub(channel)
self.event_timeout = event_timeout
Expand Down Expand Up @@ -357,7 +358,7 @@ async def create_queue(self, queue: submit_pb2.Queue) -> empty_pb2.Empty:
:param queue: A queue to create.
"""

response = await self.submit_stub.CreateQueue(queue)
response = await self.queue_stub.CreateQueue(queue)
return response

async def update_queue(self, queue: submit_pb2.Queue) -> empty_pb2.Empty:
Expand All @@ -367,7 +368,7 @@ async def update_queue(self, queue: submit_pb2.Queue) -> empty_pb2.Empty:
:param queue: A queue to update.
"""

response = await self.submit_stub.UpdateQueue(queue)
response = await self.queue_stub.UpdateQueue(queue)
return response

async def create_queues(
Expand All @@ -380,7 +381,7 @@ async def create_queues(
"""

queue_list = submit_pb2.QueueList(queues=queues)
response = await self.submit_stub.CreateQueues(queue_list)
response = await self.queue_stub.CreateQueues(queue_list)
return response

async def update_queues(
Expand All @@ -393,7 +394,7 @@ async def update_queues(
"""

queue_list = submit_pb2.QueueList(queues=queues)
response = await self.submit_stub.UpdateQueues(queue_list)
response = await self.queue_stub.UpdateQueues(queue_list)
return response

async def delete_queue(self, name: str) -> None:
Expand All @@ -405,7 +406,7 @@ async def delete_queue(self, name: str) -> None:
:return: None
"""
request = submit_pb2.QueueDeleteRequest(name=name)
await self.submit_stub.DeleteQueue(request)
await self.queue_stub.DeleteQueue(request)

async def get_queue(self, name: str) -> submit_pb2.Queue:
"""Get the queue by name.
Expand All @@ -416,7 +417,7 @@ async def get_queue(self, name: str) -> submit_pb2.Queue:
:return: A queue object. See the api definition.
"""
request = submit_pb2.QueueGetRequest(name=name)
response = await self.submit_stub.GetQueue(request)
response = await self.queue_stub.GetQueue(request)
return response

@staticmethod
Expand Down

0 comments on commit 6f38664

Please sign in to comment.