From 30a4313710b328fc197fad375e30de232c3f3e2f Mon Sep 17 00:00:00 2001 From: martha-johnston Date: Wed, 29 May 2024 12:45:38 -0400 Subject: [PATCH] add SetRPM to python SDK --- examples/server/v1/components.py | 6 +++++ src/viam/components/motor/client.py | 14 ++++++++++ src/viam/components/motor/motor.py | 25 +++++++++++++++++ src/viam/components/motor/service.py | 11 ++++++++ tests/mocks/components.py | 12 +++++++++ tests/test_motor.py | 40 ++++++++++++++++++++++++++++ 6 files changed, 108 insertions(+) diff --git a/examples/server/v1/components.py b/examples/server/v1/components.py index e4131219a..e958a0632 100644 --- a/examples/server/v1/components.py +++ b/examples/server/v1/components.py @@ -592,6 +592,12 @@ async def go_to(self, rpm: float, position_revolutions: float, extra: Optional[D self.position += rps self.powered = False + async def set_rpm(self, rpm: float, extra: Optional[Dict[str, Any]] = None, **kwargs): + if self.task: + self.task.cancel() + self.powered = True + self.task = asyncio.create_task(self.run_continuously(rpm)) + async def reset_zero_position(self, offset: float, extra: Optional[Dict[str, Any]] = None, **kwargs): if (self.position > 0 and offset > 0) or (self.position < 0 and offset < 0): self.position = offset diff --git a/src/viam/components/motor/client.py b/src/viam/components/motor/client.py index dc6e0ff14..7535bb0e7 100644 --- a/src/viam/components/motor/client.py +++ b/src/viam/components/motor/client.py @@ -17,6 +17,7 @@ MotorServiceStub, ResetZeroPositionRequest, SetPowerRequest, + SetRPMRequest, StopRequest, ) from viam.resource.rpc_client_base import ReconfigurableResourceRPCClientBase @@ -76,6 +77,19 @@ async def go_to( request = GoToRequest(name=self.name, rpm=rpm, position_revolutions=position_revolutions, extra=dict_to_struct(extra)) await self.client.GoTo(request, timeout=timeout) + async def set_rpm( + self, + rpm: float, + *, + extra: Optional[Dict[str, Any]] = None, + timeout: Optional[float] = None, + **__, + ): + if extra is None: + extra = {} + request = SetRPMRequest(name=self.name, rpm=rpm, extra=dict_to_struct(extra)) + await self.client.SetRPM(request, timeout=timeout) + async def reset_zero_position( self, offset: float, diff --git a/src/viam/components/motor/motor.py b/src/viam/components/motor/motor.py index 5fb5d793a..977969679 100644 --- a/src/viam/components/motor/motor.py +++ b/src/viam/components/motor/motor.py @@ -112,6 +112,31 @@ async def go_to( """ ... + @abc.abstractmethod + async def set_rpm( + self, + rpm: float, + *, + extra: Optional[Dict[str, Any]] = None, + timeout: Optional[float] = None, + **kwargs, + ): + """ + Spin the motor indefinitely at the specified speed, in revolutions per minute. + Positive ``rpm`` will result in forward movement and negative ``rpm`` will result in backwards movement + + :: + + my_motor = Motor.from_robot(robot=robot, name="my_motor") + + # Spin the motor at 75 RPM. + await my_motor.set_rpm(rpm=75) + + Args: + rpm (float): Speed at which the motor should rotate. + """ + ... + @abc.abstractmethod async def reset_zero_position( self, diff --git a/src/viam/components/motor/service.py b/src/viam/components/motor/service.py index 638e5be53..bc31c3df7 100644 --- a/src/viam/components/motor/service.py +++ b/src/viam/components/motor/service.py @@ -19,6 +19,8 @@ ResetZeroPositionResponse, SetPowerRequest, SetPowerResponse, + SetRPMRequest, + SetRPMResponse, StopRequest, StopResponse, ) @@ -64,6 +66,15 @@ async def GoTo(self, stream: Stream[GoToRequest, GoToResponse]) -> None: ) await stream.send_message(GoToResponse()) + async def SetRPM(self, stream: Stream[SetRPMRequest, SetRPMResponse]) -> None: + request = await stream.recv_message() + assert request is not None + name = request.name + motor = self.get_resource(name) + timeout = stream.deadline.time_remaining() if stream.deadline else None + await motor.set_rpm(request.rpm, extra=struct_to_dict(request.extra), timeout=timeout, metadata=stream.metadata) + await stream.send_message(SetRPMResponse()) + async def ResetZeroPosition(self, stream: Stream[ResetZeroPositionRequest, ResetZeroPositionResponse]) -> None: request = await stream.recv_message() assert request is not None diff --git a/tests/mocks/components.py b/tests/mocks/components.py index 82575efcb..25e0bca72 100644 --- a/tests/mocks/components.py +++ b/tests/mocks/components.py @@ -701,6 +701,18 @@ async def go_to( self.extra = extra self.timeout = timeout + async def set_rpm( + self, + rpm: float, + *, + extra: Optional[Dict[str, Any]] = None, + timeout: Optional[float] = None, + **kwargs, + ): + self.powered = rpm != 0 + self.extra = extra + self.timeout = timeout + async def reset_zero_position( self, offset: float, diff --git a/tests/test_motor.py b/tests/test_motor.py index e9a9c825f..8f6ed9ec7 100644 --- a/tests/test_motor.py +++ b/tests/test_motor.py @@ -19,6 +19,7 @@ MotorServiceStub, ResetZeroPositionRequest, SetPowerRequest, + SetRPMRequest, StopRequest, ) from viam.resource.manager import ResourceManager @@ -83,6 +84,19 @@ async def test_go_to(self, motor: MockMotor): await motor.go_to(-10, 50) assert await motor.get_position() == 50 + @pytest.mark.asyncio + async def test_set_rpm(self, motor: MockMotor): + await motor.set_rpm(30, timeout=4.56) + assert motor.timeout == loose_approx(4.56) + moving, pwr = await motor.is_powered() + assert moving is True + assert pwr > 0 + + await motor.set_rpm(-10) + moving, pwr = await motor.is_powered() + assert moving is True + assert pwr < 0 + @pytest.mark.asyncio async def test_reset_zero(self, motor: MockMotor): await motor.reset_zero_position(20, timeout=5.67) @@ -205,6 +219,20 @@ async def test_go_to(self, motor: MockMotor, service: MotorRPCService): await client.GoTo(request) assert motor.position == 50 + @pytest.mark.asyncio + async def test_set_rpm(self, motor: MockMotor, service: MotorRPCService): + async with ChannelFor([service]) as channel: + client = MotorServiceStub(channel) + + request = SetRPMRequest(name=motor.name, rpm=30) + await client.SetRPM(request, timeout=4.56) + assert motor.power > 0 + assert motor.timeout == loose_approx(4.56) + + request = SetRPMRequest(name=motor.name, rpm=-10) + await client.SetRPM(request) + assert motor.power < 0 + @pytest.mark.asyncio async def test_reset_zero(self, motor: MockMotor, service: MotorRPCService): async with ChannelFor([service]) as channel: @@ -340,6 +368,18 @@ async def test_go_to(self, motor: MockMotor, service: MotorRPCService): await client.go_to(-10, 50) assert motor.position == 50 + @pytest.mark.asyncio + async def test_set_rpm(self, motor: MockMotor, service: MotorRPCService): + async with ChannelFor([service]) as channel: + client = MotorClient(motor.name, channel) + + await client.set_rpm(30, timeout=4.56) + assert motor.power > 0 + assert motor.timeout == loose_approx(4.56) + + await client.set_rpm(-10) + assert motor.power < 0 + @pytest.mark.asyncio async def test_reset_zero(self, motor: MockMotor, service: MotorRPCService): async with ChannelFor([service]) as channel: