import io
import json
import unittest
from unittest.mock import patch

from as213905 import AS213905, APIError


class Response(io.BytesIO):
    def __enter__(self): return self
    def __exit__(self, *args): return False


class ClientTests(unittest.TestCase):
    def test_organization_sets_bearer_header(self):
        def open_(request, timeout):
            self.assertEqual(request.get_header("Authorization"), "Bearer secret")
            return Response(json.dumps({"organization": {"asn": 213905}}).encode())
        with patch("as213905.client.urlopen", open_):
            self.assertEqual(AS213905("secret").organization()["asn"], 213905)

    def test_traffic_query(self):
        client = AS213905("secret")
        self.assertEqual(client._traffic_query(7, None, None), "?days=7")


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