Repository: commaai/openpilot
File: system/webrtc/tests/test_stream_session.py (Python)

Code before:

      mocked_pubmaster.send.assert_called_once()
      mt, md = mocked_pubmaster.send.call_args.args
      assert mt == msg["type"]
      assert isinstance(md, capnp._DynamicStructBuilder)
      assert hasattr(md, msg["type"])

      mocked_pubmaster.reset_mock()

  def test_livestream_track(self, mocker):
    fake_msg = messaging.new_message("livestreamDriverEncodeData")

    config = {"receive.return_value": fake_msg.to_bytes()}
    mocker.patch("msgq.SubSocket", spec=True, **config)
    track = LiveStreamVideoStreamTrack("driver")

    assert track.id.startswith("driver")
    assert track.codec_preference() == "H264"

    for i in range(5):
      packet = self.loop.run_until_complete(track.recv())
      assert packet.time_base == VIDEO_TIME_BASE
      if i == 0:
        start_ns = time.monotonic_ns()
        start_pts = packet.pts
      assert abs(i+ packet.pts - (start_pts + (((time.monotonic_ns()-start_ns) * VIDEO_CLOCK_RATE)//1_000_000_000))) < 450 #5ms
      assert packet.size == 0

  def test_input_audio_track(self, mocker):
    packet_time, rate = 0.02, 16000
    sample_count = int(packet_time * rate)
    mocked_stream = mocker.MagicMock(spec=pyaudio.Stream)
    mocked_stream.read.return_value = b"\x00" * 2 * sample_count

    config = {"open.side_effect": lambda *args, **kwargs: mocked_stream}
    mocker.patch("pyaudio.PyAudio", spec=True, **config)
    track = AudioInputStreamTrack(audio_format=pyaudio.paInt16, packet_time=packet_time, rate=rate)

    for i in range(5):
      frame = self.loop.run_until_complete(track.recv())
      assert frame.rate == rate
      assert frame.samples == sample_count
      assert frame.pts == i * sample_count


Review comment (suggestion):
```suggestion
      assert abs(i + packet.pts - (start_pts + (((time.monotonic_ns() - start_ns) * VIDEO_CLOCK_RATE) // 1_000_000_000))) < 450 #5ms
```

Code after:

      mocked_pubmaster.send.assert_called_once()
      mt, md = mocked_pubmaster.send.call_args.args
      assert mt == msg["type"]
      assert isinstance(md, capnp._DynamicStructBuilder)
      assert hasattr(md, msg["type"])

      mocked_pubmaster.reset_mock()

  def test_livestream_track(self, mocker):
    fake_msg = messaging.new_message("livestreamDriverEncodeData")

    config = {"receive.return_value": fake_msg.to_bytes()}
    mocker.patch("msgq.SubSocket", spec=True, **config)
    track = LiveStreamVideoStreamTrack("driver")

    assert track.id.startswith("driver")
    assert track.codec_preference() == "H264"

    for i in range(5):
      packet = self.loop.run_until_complete(track.recv())
      assert packet.time_base == VIDEO_TIME_BASE
      if i == 0:
        start_ns = time.monotonic_ns()
        start_pts = packet.pts
      assert abs(i + packet.pts - (start_pts + (((time.monotonic_ns() - start_ns) * VIDEO_CLOCK_RATE) // 1_000_000_000))) < 450 #5ms
      assert packet.size == 0

  def test_input_audio_track(self, mocker):
    packet_time, rate = 0.02, 16000
    sample_count = int(packet_time * rate)
    mocked_stream = mocker.MagicMock(spec=pyaudio.Stream)
    mocked_stream.read.return_value = b"\x00" * 2 * sample_count

    config = {"open.side_effect": lambda *args, **kwargs: mocked_stream}
    mocker.patch("pyaudio.PyAudio", spec=True, **config)
    track = AudioInputStreamTrack(audio_format=pyaudio.paInt16, packet_time=packet_time, rate=rate)

    for i in range(5):
      frame = self.loop.run_until_complete(track.recv())
      assert frame.rate == rate
      assert frame.samples == sample_count
      assert frame.pts == i * sample_count