"""Holding a connection open across several requests.

The reply path has two shapes. A small answer goes out in one write and the
connection is done with immediately. A large one is queued and finishes on a
later pass through the event loop, by which time the workstruct that served
the request has moved on, so whether to keep the connection has to be
remembered on the connection itself.
"""

import unittest
import urllib.parse

from harness import TrackerTestCase, bdecode


class KeepAlive(TrackerTestCase):
    def announce_path(self, name, port=6881):
        return "/announce?info_hash=%s&peer_id=%s&port=%d&left=0" % (
            urllib.parse.quote_from_bytes(self.hash_for(name), safe=""),
            "x" * 20,
            port,
        )

    def test_two_requests_over_one_connection(self):
        c = self.tracker.connect()
        try:
            c.send(self.announce_path("kaA"))
            status, _, first = c.reply()
            self.assertIn(b"200", status)
            c.send(self.announce_path("kaB"))
            status, _, second = c.reply()
            self.assertIn(b"200", status)
            self.assertIn(b"complete", bdecode(second))
        finally:
            c.close()

    def test_close_ends_the_connection(self):
        c = self.tracker.connect()
        try:
            c.send(self.announce_path("kaclose"), connection="close")
            c.reply()
            self.assertFalse(c.is_open(), "the tracker was asked to hang up")
        finally:
            c.close()

    def test_a_connection_survives_a_reply_too_large_for_one_write(self):
        """The case that was broken: a queued reply finished draining and the
        connection was closed regardless of what the client had asked for."""
        for i in range(400):
            self.tracker.announce(bytes([i // 256, i % 256]) + b"\0" * 18, port=7000 + i, left=0)

        c = self.tracker.connect()
        try:
            c.send("/scrape")
            status, headers, body = c.reply()
            self.assertIn(b"200", status)
            self.assertGreater(len(body), 16384, "this reply has to outgrow one write")

            c.send(self.announce_path("kaafter"))
            status, _, after = c.reply()
            self.assertIn(b"200", status)
            self.assertIn(b"complete", bdecode(after))
        finally:
            c.close()

    def test_several_requests_in_a_row(self):
        c = self.tracker.connect()
        try:
            for i in range(20):
                c.send(self.announce_path("karow%d" % i))
                status, _, body = c.reply()
                self.assertIn(b"200", status)
            self.assertEqual(self.tracker.torrent_count(), 20)
        finally:
            c.close()

    def test_without_the_header_the_connection_is_not_held(self):
        """HTTP/1.1 calls a connection persistent unless told otherwise; this
        tracker holds one only when asked by name. Pinned as it is."""
        c = self.tracker.connect()
        try:
            c.send(self.announce_path("kanone"), connection=None)
            c.reply()
            self.assertFalse(c.is_open())
        finally:
            c.close()

    def test_the_error_path_does_not_hold_the_connection(self):
        c = self.tracker.connect()
        try:
            c.send("/announce?info_hash=short")
            status, _, _ = c.reply()
            self.assertIn(b"400", status)
        finally:
            c.close()

    def test_an_error_after_a_kept_connection_is_still_an_error(self):
        """The keep-alive decision used to be taken after the handler had run,
        so an error reply inherited whatever the previous request wanted."""
        c = self.tracker.connect()
        try:
            c.send(self.announce_path("kabefore"))
            self.assertIn(b"200", c.reply()[0])
            c.send("/announce?info_hash=short")
            self.assertIn(b"400", c.reply()[0])
        finally:
            c.close()


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