[Bug]: Python data channel masks saved reader exception with Channel closed prematurely #40774

import unittest

from apache_beam.runners.worker import data_plane


class ChannelErrorTest(unittest.TestCase):
    def test_saved_reader_failure_is_preserved(self):
        channel = data_plane._GrpcDataChannel()
        error = ValueError("original read failure")
        self.addCleanup(self.assertRaises, ValueError, channel.close)

        def failed_inputs():
            raise error
            yield

        with self.assertLogs(data_plane._LOGGER, level="ERROR"):
            with self.assertRaises(ValueError):
                channel._read_inputs(failed_inputs())
        with self.assertRaisesRegex(
                RuntimeError, "Channel closed prematurely") as caught:
            list(channel.input_elements("instruction", ["input"]))
        self.assertIs(caught.exception.__cause__, error)


if __name__ == "__main__":
    unittest.main()

if self._closed:
    raise RuntimeError(
        'Channel closed prematurely.') from self._exception