Robotic automation factory

"""
 Few comments:
 - assuming integers only for simplification
 - I'd use pydantic for checking inputs
 - I'd use pytest for real testing
 - I'd use pytest to test if exceptions were raised
"""


def sort(width: int, height: int, length: int, mass: int) -> str:
  """
  Determine how a package should be dispatched based on its dimensions and mass.
  Returns one of: "STANDARD", "SPECIAL", or "REJECTED".
  """

  for arg in (width, height, length, mass):
    if not isinstance(arg, int):
      raise TypeError("All inputs must be integers") # I'd use pytest to catch this exception
    if arg < 0:
      raise ValueError("All dimensions and mass must be non-negative") # I'd use pytest to catch this exception

  # Analyze packet:
  bulky = False
  heavy = False

  volume = width * height * length
  cond_volume = volume >= 1_000_000

  cond_dimension = (width >= 150) or (height >= 150) or (length >= 150)

  if cond_volume or cond_dimension:
    bulky = True

  if mass >= 20:
    heavy = True

  # Select dispatch
  if bulky ^ heavy:
    return "SPECIAL"
  elif bulky and heavy:
    return "REJECTED"
  else:
    return "STANDARD"

  return "ERROR" # I leave this here in case I mess up sth while coding.



def main():
  print("Sorting...")

  # standard packet:
  assert(sort(1, 2, 3, 4) == "STANDARD")

  # bulky because of dimension and volume
  assert(sort(100,200,300, 4) == "SPECIAL")

  # bulky because of dimension(width)
  assert (sort(149, 2, 3, 4) == "STANDARD")
  assert (sort(150, 2, 3, 4) == "SPECIAL")
  assert (sort(151, 2, 3, 4) == "SPECIAL")

  # bulky because of dimension(height)
  assert (sort(10, 149, 3, 4) == "STANDARD")
  assert (sort(10, 150, 3, 4) == "SPECIAL")

  # bulky because of dimension(length)
  assert (sort(10, 20, 149, 4) == "STANDARD")
  assert(sort(10,20,150, 4) == "SPECIAL")

  # bulky because of volume
  assert (sort(100, 100, 99, 4) == "STANDARD")
  assert (sort(100, 100, 100, 4) == "SPECIAL")

  # bulky and heavy
  assert (sort(100, 200, 300, 25) == "REJECTED")

  # SPECIAL because it's heavy
  assert(sort(1,2,3, 19) == "STANDARD")
  assert (sort(1, 2, 3, 20) == "SPECIAL")

  # testing wrong inputs:
  try:
    sort(-1, 2, 3, 4)
    assert False
  except ValueError:
    pass

  try:
    sort(1, 1.2, 3, 4)
    assert False
  except TypeError:
    pass

  try:
    sort(1, 2, "3", 4)
    assert False
  except TypeError:
    pass

  print("All tests passed!")


if __name__ == "__main__":
  main()