def __init__(self, bike_connection_manager: BikeConnectionManager, rental_manager: RentalManager): self.reservations: Dict[int, Set[Tuple[int, int, datetime]]] = defaultdict(set) """Maps a pickup point to a list of reservations (and user id, and date for).""" self.pickup_points: Set[PickupPoint] = set() self._bike_connection_manager = bike_connection_manager self._rental_manager = rental_manager self.hub = EventHub(ReservationEvent)
async def test_trigger_event_natural_syntax(self): """Assert that events can be triggered with the natural syntax.""" def raise_listener(argument): raise TestException("This runs!") emitter = EventHub(ExampleEvents) emitter.something_happened += raise_listener with pytest.raises(TestException): emitter.something_happened("test")
async def test_trigger_event(self): emitter = EventHub(ExampleEvents) def raise_listener(argument): raise TestException("This runs!") with pytest.raises(TestException): emitter.subscribe(ExampleEvents.something_happened, raise_listener) emitter.emit(ExampleEvents.something_happened, argument="test")
class ReservationManager(Rebuildable): def __init__(self, bike_connection_manager: BikeConnectionManager, rental_manager: RentalManager): self.reservations: Dict[int, Set[Tuple[int, int, datetime]]] = defaultdict(set) """Maps a pickup point to a list of reservations (and user id, and date for).""" self.pickup_points: Set[PickupPoint] = set() self._bike_connection_manager = bike_connection_manager self._rental_manager = rental_manager self.hub = EventHub(ReservationEvent) async def reserve(self, user: User, pickup: PickupPoint, for_time: datetime) -> Reservation: """ Reserves a bike at a pickup point. :raises ReservationError: If there are any issues with making the reservation. """ if for_time - datetime.now(timezone.utc) < MINIMUM_RESERVATION_TIME: # ensure there is a bike there available_bikes = await self.available_bikes(pickup) if len(available_bikes) <= len(self.reservations[pickup.id]): raise ReservationError( f"No available bikes in the requested point, and not enough time to source one (less than {MINIMUM_RESERVATION_TIME}).") reservation = await create_reservation(user, pickup, for_time) self.reservations[pickup.id].add((reservation.id, user.id, reservation.reserved_for)) self.hub.opened_reservation(pickup, user, for_time) return reservation async def claim(self, user: User, bike: Bike) -> Tuple[Rental, Point]: """ Collects the bike for a given reservation, creating a rental. Collection can only happen 30 minutes before or after :param user: The user collecting the reservation. :param bike: The specific bike to collect. :raises CollectionError: If the bike is being picked up outside the reservation window. :raises ReservationError: If there are no bikes in the pickup point. :raises ActiveRentalError: If the user already has an active rental. :raises ValueError: If the pickup points """ collection_location = self._bike_connection_manager.most_recent_location(bike) if collection_location is None: raise CollectionError("Could not find bike.") # THIS SHOULD NEVER HAPPEN else: collection_location, _, _ = collection_location active_reservations = await current_reservations(user) valid_reservations = [] for reservation in active_reservations: upper_bound = reservation.reserved_for + RESERVATION_WINDOW / 2 if datetime.now(timezone.utc) >= upper_bound: reservation.outcome = ReservationOutcome.EXPIRED await reservation.save() elif collection_location.within(reservation.pickup_point.area): valid_reservations.append(reservation) if not valid_reservations: raise CollectionError("User has no active reservations at the pickup point of the requested bike.") reservation = sorted(valid_reservations, key=lambda x: x.reserved_for)[0] lower_bound = reservation.reserved_for - RESERVATION_WINDOW / 2 upper_bound = reservation.reserved_for + RESERVATION_WINDOW / 2 if not lower_bound <= datetime.now(timezone.utc) <= upper_bound: raise CollectionError( f"You may only collect a rental in a {RESERVATION_WINDOW} window around {reservation.reserved_for}." ) available_bikes = await self.available_bikes(reservation.pickup_point) if not available_bikes: raise ReservationError("No bikes at this pickup point.") pickup = self._pickup_containing(bike) if not pickup or pickup.id != reservation.pickup_point.id: raise CollectionError("Requested bike is not in the pickup point of the reservation.") if bike not in available_bikes: raise CurrentlyRentedError("Requested bike is currently being rented.", available_bikes) rental, start_location = await self._rental_manager.create(reservation.user, available_bikes[0]) reservation.claimed_rental = rental await self._close_reservation(reservation, ReservationOutcome.CANCELLED) return rental, start_location async def cancel(self, reservation: Reservation): """ Cancels a reservation. """ await self._close_reservation(reservation, ReservationOutcome.CANCELLED) self.hub.cancelled_reservation(reservation.pickup_point, reservation.user, reservation.reserved_for) def reservations_in(self, pickup_point: int) -> List[int]: """Gets the reservation ids for a given pickup point.""" return [rid for rid, uid, time in self.reservations[pickup_point]] def is_reserved(self, bike: Bike) -> bool: """Checks if the given bike is reserved.""" try: pickup = self._pickup_containing(bike) except ValueError: # the bike has no location updates. it can't be reserved return False if pickup is None: return False bikes = self._bike_connection_manager.bikes_in(pickup.area) return len(bikes) <= len(self.reservations[pickup.id]) def pickup_bike_surplus(self, pickup_point) -> int: """ Returns how many more free bikes there are than reservations. A negative value indicates a shortage. """ available_bike_count = len(self._bike_connection_manager.bikes_in(pickup_point.area)) bikes_needed_for_reservations = len([ rid for rid, uid, time in self.reservations[pickup_point.id] if time < datetime.now(timezone.utc) + MINIMUM_RESERVATION_TIME ]) return available_bike_count - bikes_needed_for_reservations async def _rebuild(self): """Rebuilds the reservation manager from the database.""" unhandled_reservations = await Reservation.filter(outcome__isnull=True).prefetch_related("pickup_point", "user") for reservation in unhandled_reservations: self.reservations[reservation.pickup_point.id].add( (reservation.id, reservation.user_id, reservation.reserved_for)) self.pickup_points.add(reservation.pickup_point) self.hub.opened_reservation(reservation.pickup_point, reservation.user, reservation.reserved_for) async def _close_reservation(self, reservation: Reservation, outcome: ReservationOutcome): self.reservations[reservation.pickup_point_id].remove( (reservation.id, reservation.user_id, reservation.reserved_for)) reservation.ended_at = datetime.now(timezone.utc) reservation.outcome = outcome await reservation.save() async def available_bikes(self, pickup_point: PickupPoint) -> List[Bike]: """Gets the available bikes in a pickup point.""" bike_ids = self._bike_connection_manager.bikes_in(pickup_point.area) return await self._rental_manager.get_available_bikes_out_of(bike_ids) def _pickup_containing(self, bike: Bike) -> Optional[PickupPoint]: """Gets the pickup point the bike is currently in.""" if not self._bike_connection_manager.is_connected(bike): raise ValueError("Bike not connected!") bike_location = self._bike_connection_manager.most_recent_location(bike) if bike_location is None: raise ValueError("Bike has not submitted its location yet!") else: _, _, pickup = bike_location return pickup
def __init__(self, payment_manager: PaymentManager): self._active_rentals: Dict[int, Tuple[int, int]] = {} """Maps user ids to a tuple containing their current rental and current bike.""" self.hub = EventHub(RentalEvent) self._payment_manager = payment_manager
class RentalManager(Rebuildable): """ Handles the lifecycle of the rental in the system. Also publishes events on its hub, so that other modules can stay up to date with the system. When the module starts, it will replay all events from midnight that night. """ def __init__(self, payment_manager: PaymentManager): self._active_rentals: Dict[int, Tuple[int, int]] = {} """Maps user ids to a tuple containing their current rental and current bike.""" self.hub = EventHub(RentalEvent) self._payment_manager = payment_manager async def create(self, user: User, bike: Bike) -> Tuple[Rental, Point]: """ Creates a new rental for a user. :raises ActiveRentalError: If the requested user currently has a rental active. :raises CurrentlyRentedError: If the requested bike is currently in use. """ if user.id in self._active_rentals: raise ActiveRentalError(self._active_rentals[user.id]) if self.is_in_use(bike): raise CurrentlyRentedError("The requested bike is in use.", []) if not await self._payment_manager.is_customer(user): raise CustomerError() rental = await Rental.create(user=user, bike=bike) rental.bike = bike rental.user = user await self._publish_event(rental, RentalUpdateType.RENT) self._active_rentals[user.id] = (rental.id, bike.id) await rental.fetch_related("updates") current_location = await LocationUpdate.filter(bike=bike).first() self.hub.emit(RentalEvent.rental_started, user, bike, current_location.location) return rental, current_location.location if current_location is not None else None async def finish(self, user: User, *, extra_cost=0.0) -> Tuple[Rental, str]: """ Completes a rental, charging the customer. :returns: A tuple with the rental and the URL of the receipt. :raises InactiveRentalError: When there is no active rental for that user, or the given rental is not active. """ rental, distance = await self._get_rental_with_distance(user) rental.price = await get_price(rental.updates[0].time, datetime.now(), extra_cost) current_location = await LocationUpdate.filter(bike_id=rental.bike_id ).first() success, receipt_url = await self._payment_manager.charge_customer( user, rental, distance) if success: del self._active_rentals[user.id] await rental.save() await self._publish_event(rental, RentalUpdateType.RETURN) self.hub.emit(RentalEvent.rental_ended, user, rental.bike, current_location.location, rental.price, distance) else: raise Exception("Couldn't charge card!") return rental, receipt_url async def cancel(self, user: User) -> Rental: """Cancels a rental, effective immediately, waiving the rental fee.""" rental, distance = await self._get_rental_with_distance(user) await self._publish_event(rental, RentalUpdateType.CANCEL) self.hub.emit(RentalEvent.rental_cancelled, user, rental.bike) del self._active_rentals[user.id] return rental async def active_rentals(self) -> List[Rental]: """Gets all the active rentals.""" active_rental_ids = [ rental_id for rental_id, bike_id in self._active_rentals.values() ] return await Rental.filter(id__in=active_rental_ids ).prefetch_related('updates', 'bike') async def active_rental( self, user: Union[User, int], *, with_locations=False) -> Optional[Union[Rental, Tuple]]: """Gets the active rental for a given user.""" if isinstance(user, int): user_id = user if isinstance(user, User): user_id = user.id rental_id, bike_id = self._active_rentals.get(user_id, (None, None)) if rental_id is None: return None rental = await Rental.filter(id=rental_id).first().prefetch_related( 'updates', 'bike') if with_locations: locations = await LocationUpdate.filter( bike_id=bike_id, time__gte=rental.updates[0].time.strftime("%Y-%m-%d %H:%M:%S")) if locations: return rental, locations[0].location, locations[-1].location else: return rental, None, None else: return rental def has_active_rental(self, user: Union[User, int]) -> bool: """Checks if the given user has an active rental.""" return (user if isinstance(user, int) else user.id) in self._active_rentals def is_active(self, rental_id): """Checks if the given rental ID is currently active.""" return rental_id in { rental for rental, bike in self._active_rentals.values() } def is_in_use(self, bike: Union[Bike, int]) -> bool: """Checks if the given bike is in use.""" bid = bike.id if isinstance(bike, Bike) else bike return any(bid == rental[1] for rental in self._active_rentals.values()) def is_available(self, bike: Union[Bike, int], reservation_manager) -> bool: """A bike is available if the bike is un-rented and it is not reserved.""" return not self.is_in_use( bike) and not reservation_manager.is_reserved(bike) def is_renting(self, user_id: int, bike_id: int) -> bool: """Checks if the given user is renting the given bike.""" return user_id in self._active_rentals and self._active_rentals[ user_id][1] == bike_id async def get_price_estimate(self, rental: Union[Rental, int]) -> float: """Gets the price of the rental so far.""" if isinstance(rental, int): rental = await Rental.filter(id=rental).first() elif isinstance(rental, Rental): rental = rental return await get_price(rental.start_time, datetime.now()) async def get_available_bikes_out_of(self, bike_ids: List[int]) -> List[Bike]: """Given a list of bike ids, checks if they are free or not and returns the ones that are free.""" used_bikes = { bike_id for rental_id, bike_id in self._active_rentals.values() } available_bikes = set(bike_ids) - used_bikes if not available_bikes: return [] query = Bike.filter(id__in=available_bikes) return await query.prefetch_related( Prefetch("location_updates", queryset=LocationUpdate.all().limit(100)), "state_updates", Prefetch("issues", queryset=Issue.filter(status__not=IssueStatus.CLOSED))) async def _rebuild(self): """ Rebuilds the currently active rentals from the database. Also replays events that happened today for use by subscribers. """ async for rental in await unfinished_rentals(): self._active_rentals[rental.user_id] = (rental.id, rental.bike_id) midnight = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0) todays_updates = await RentalUpdate.filter(time__gt=midnight ).prefetch_related( "rental", "rental__user", "rental__bike") for update in todays_updates: if update.type == RentalUpdateType.RENT: self.hub.emit(RentalEvent.rental_started, update.rental.user, update.rental.bike, None) # todo location elif update.type == RentalUpdateType.RETURN: self.hub.emit(RentalEvent.rental_ended, update.rental.user, update.rental.bike, None, update.rental.price, None) elif update.type == RentalUpdateType.CANCEL: self.hub.emit(RentalEvent.rental_cancelled, update.rental.user, update.rental.bike) async def _get_rental_with_distance(self, user: User) -> Tuple[Rental, float]: """Given a rental or user, "resolves" the rental, user pair.""" if not isinstance(user, User): raise TypeError( f"Supplied target must be a Rental or User, not {type(user)}") if user.id not in self._active_rentals: raise InactiveRentalError("Given user has no active rentals!") rental_id, bike_id = self._active_rentals[user.id] rental, distance = await get_rental_with_distance(rental_id) return rental, distance @staticmethod async def _publish_event(rental: Rental, event_type: RentalUpdateType): update = await RentalUpdate.create(rental=rental, type=event_type) if rental.updates._fetched: rental.updates.related_objects.append(update)
async def test_bad_unsubscribe(self): """Assert that unsubscribing a handler that isn't registered fails.""" emitter = EventHub() with pytest.raises(NoSuchListenerError): emitter.unsubscribe(ExampleEvents.something_happened, self.handler)
async def test_unsubscribe_from_event(self): """Assert that a handler can be unsubscribed from an event.""" emitter = EventHub(ExampleEvents) emitter.subscribe(emitter.something_happened, self.handler) emitter.unsubscribe(emitter.something_happened, self.handler)
async def test_subscribe_to_event_through_emitter(self): """Assert that an event can also be referenced through the emitter.""" emitter = EventHub(ExampleEvents) assert sum((len(l) for l in emitter._listeners.values()), 0) == 0 emitter.subscribe(emitter.something_happened, self.handler) assert sum((len(l) for l in emitter._listeners.values()), 0) != 0
async def test_subscriber_has_similar_signature(self): """Assert that a subscriber to an event must have a similar signature.""" emitter = EventHub(ExampleEvents) with pytest.raises(InvalidHandlerError): emitter.subscribe(ExampleEvents.something_happened, self.invalid_handler)
async def test_subscribe_to_event(self): """Assert that a handler can be registered on a hub's event.""" hub = EventHub(ExampleEvents) assert sum((len(l) for l in hub._listeners.values()), 0) == 0 hub.subscribe(ExampleEvents.something_happened, self.handler) assert sum((len(l) for l in hub._listeners.values()), 0) != 0
async def test_missing_event_on_emitter(self): """Assert that getting a non-existent event on an emitter raises an error.""" emitter = EventHub() with pytest.raises(NoSuchEventError): emitter.bad_event
async def test_event_on_emitter(self): """Assert that an event can be accessed through the hub.""" hub = EventHub(ExampleEvents) assert hub.something_happened.event == ExampleEvents.something_happened
async def test_event_in_emitter(self): """Assert that you can check the existence of an event on an hub.""" hub = EventHub(ExampleEvents) assert ExampleEvents.something_happened in hub assert SecondExampleEvents.something_else_happened not in hub
async def test_event_list_in_emitter(self): """Assert that you can check the existence of an event list on a hub.""" hub = EventHub(ExampleEvents) assert ExampleEvents in hub assert SecondExampleEvents not in hub
async def test_unsubscribe_to_event_natural_syntax(self): emitter = EventHub(ExampleEvents) emitter.subscribe(ExampleEvents.something_happened, self.handler) emitter.something_happened -= self.handler