#!/usr/bin/env vpython
# Copyright (c) 2012 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# [VPYTHON:BEGIN]
# wheel: <
#   name: "infra/python/wheels/pytz-py2_py3"
#   version: "version:2018.4"
# >
# wheel: <
#   name: "infra/python/wheels/tempora-py2_py3"
#   version: "version:1.11"
# >
# wheel: <
#   name: "infra/python/wheels/more-itertools-py2_py3"
#   version: "version:4.1.0"
# >
# wheel: <
#   name: "infra/python/wheels/backports_functools_lru_cache-py2_py3"
#   version: "version:1.5"
# >
# wheel: <
#   name: "infra/python/wheels/six-py2_py3"
#   version: "version:1.12.0"
# >
# wheel: <
#   name: "infra/python/wheels/portend-py2_py3"
#   version: "version:2.2"
# >
# wheel: <
#   name: "infra/python/wheels/cheroot-py2_py3"
#   version: "version:6.2.4"
# >
# wheel: <
#   name: "infra/python/wheels/cherrypy-py2_py3"
#   version: "version:14.2.0"
# >
# [VPYTHON:END]

"""Tests for Constrained Network Server."""
import os
import signal
import subprocess
import tempfile
import time
import unittest
import urllib2
import cherrypy
import cns
import traffic_control

# The local interface to test on.
_INTERFACE = 'lo'


class PortAllocatorTest(unittest.TestCase):
  """Unit tests for the Port Allocator class."""

  # Expiration time for ports. In mock time.
  _EXPIRY_TIME = 6

  def setUp(self):
    # Mock out time.time() to accelerate port expiration testing.
    self._old_time = time.time
    self._current_time = 0
    time.time = lambda: self._current_time

    # TODO(dalecurtis): Mock out actual calls to shadi's port setup.
    self._pa = cns.PortAllocator(cns._DEFAULT_CNS_PORT_RANGE, self._EXPIRY_TIME)
    self._MockTrafficControl()

  def tearDown(self):
    self._pa.Cleanup(all_ports=True)
    # Ensure ports are cleaned properly.
    self.assertEquals(self._pa._ports, {})
    time.time = self._old_time
    self._RestoreTrafficControl()

  def _MockTrafficControl(self):
    self.old_CreateConstrainedPort = traffic_control.CreateConstrainedPort
    self.old_DeleteConstrainedPort = traffic_control.DeleteConstrainedPort
    self.old_TearDown = traffic_control.TearDown

    traffic_control.CreateConstrainedPort = lambda config: True
    traffic_control.DeleteConstrainedPort = lambda config: True
    traffic_control.TearDown = lambda config: True

  def _RestoreTrafficControl(self):
    traffic_control.CreateConstrainedPort = self.old_CreateConstrainedPort
    traffic_control.DeleteConstrainedPort = self.old_DeleteConstrainedPort
    traffic_control.TearDown = self.old_TearDown

  def testPortAllocator(self):
    # Ensure Get() succeeds and returns the correct port.
    self.assertEquals(self._pa.Get('test'), cns._DEFAULT_CNS_PORT_RANGE[0])

    # Call again with the same key and make sure we get the same port.
    self.assertEquals(self._pa.Get('test'), cns._DEFAULT_CNS_PORT_RANGE[0])

    # Call with a different key and make sure we get a different port.
    self.assertEquals(self._pa.Get('test2'), cns._DEFAULT_CNS_PORT_RANGE[0] + 1)

    # Update fake time so that ports should expire.
    self._current_time += self._EXPIRY_TIME + 1

    # Test to make sure cache is checked before expiring ports.
    self.assertEquals(self._pa.Get('test2'), cns._DEFAULT_CNS_PORT_RANGE[0] + 1)

    # Update fake time so that ports should expire.
    self._current_time += self._EXPIRY_TIME + 1

    # Request a new port, old ports should be expired, so we should get the
    # first port in the range. Make sure this is the only allocated port.
    self.assertEquals(self._pa.Get('test3'), cns._DEFAULT_CNS_PORT_RANGE[0])
    self.assertEquals(self._pa._ports.keys(), [cns._DEFAULT_CNS_PORT_RANGE[0]])

  def testPortAllocatorExpiresOnlyCorrectPorts(self):
    # Ensure Get() succeeds and returns the correct port.
    self.assertEquals(self._pa.Get('test'), cns._DEFAULT_CNS_PORT_RANGE[0])

    # Stagger port allocation and so we can ensure only ports older than the
    # expiry time are actually expired.
    self._current_time += self._EXPIRY_TIME / 2 + 1

    # Call with a different key and make sure we get a different port.
    self.assertEquals(self._pa.Get('test2'), cns._DEFAULT_CNS_PORT_RANGE[0] + 1)

    # After this sleep the port with key 'test' should expire on the next Get().
    self._current_time += self._EXPIRY_TIME / 2 + 1

    # Call with a different key and make sure we get the first port.
    self.assertEquals(self._pa.Get('test3'), cns._DEFAULT_CNS_PORT_RANGE[0])
    self.assertEquals(set(self._pa._ports.keys()), set([
        cns._DEFAULT_CNS_PORT_RANGE[0], cns._DEFAULT_CNS_PORT_RANGE[0] + 1]))

  def testPortAllocatorNoExpiration(self):
    # Setup PortAllocator w/o port expiration.
    self._pa = cns.PortAllocator(cns._DEFAULT_CNS_PORT_RANGE, 0)

    # Ensure Get() succeeds and returns the correct port.
    self.assertEquals(self._pa.Get('test'), cns._DEFAULT_CNS_PORT_RANGE[0])

    # Update fake time to see if ports expire.
    self._current_time += self._EXPIRY_TIME

    # Send second Get() which would normally cause ports to expire. Ensure that
    # the ports did not expire.
    self.assertEquals(self._pa.Get('test2'), cns._DEFAULT_CNS_PORT_RANGE[0] + 1)
    self.assertEquals(set(self._pa._ports.keys()), set([
        cns._DEFAULT_CNS_PORT_RANGE[0], cns._DEFAULT_CNS_PORT_RANGE[0] + 1]))

  def testPortAllocatorCleanMatchingIP(self):
    # Setup PortAllocator w/o port expiration.
    self._pa = cns.PortAllocator(cns._DEFAULT_CNS_PORT_RANGE, 0)

    # Ensure Get() succeeds and returns the correct port.
    self.assertEquals(self._pa.Get('ip1', t=1), cns._DEFAULT_CNS_PORT_RANGE[0])
    self.assertEquals(self._pa.Get('ip1', t=2),
                      cns._DEFAULT_CNS_PORT_RANGE[0] + 1)
    self.assertEquals(self._pa.Get('ip1', t=3),
                      cns._DEFAULT_CNS_PORT_RANGE[0] + 2)
    self.assertEquals(self._pa.Get('ip2', t=1),
                      cns._DEFAULT_CNS_PORT_RANGE[0] + 3)

    self._pa.Cleanup(all_ports=False, request_ip='ip1')

    self.assertEquals(self._pa._ports.keys(),
                      [cns._DEFAULT_CNS_PORT_RANGE[0] + 3])
    self.assertEquals(self._pa.Get('ip2'), cns._DEFAULT_CNS_PORT_RANGE[0])
    self.assertEquals(self._pa.Get('ip1'), cns._DEFAULT_CNS_PORT_RANGE[0] + 1)

    self._pa.Cleanup(all_ports=False, request_ip='ip2')
    self.assertEquals(self._pa._ports.keys(),
                      [cns._DEFAULT_CNS_PORT_RANGE[0] + 1])

    self._pa.Cleanup(all_ports=False, request_ip='abc')
    self.assertEquals(self._pa._ports.keys(),
                      [cns._DEFAULT_CNS_PORT_RANGE[0] + 1])

    self._pa.Cleanup(all_ports=False, request_ip='ip1')
    self.assertEquals(self._pa._ports.keys(), [])


class ConstrainedNetworkServerTest(unittest.TestCase):
  """End to end tests for ConstrainedNetworkServer system.

  These tests require root access and run the cherrypy server along with
  tc/iptables commands.
  """

  # Amount of time to wait for the CNS to start up.
  _SERVER_START_SLEEP_SECS = 1

  # Sample data used to verify file serving.
  _TEST_DATA = 'The quick brown fox jumps over the lazy dog'

  # Server information.
  _SERVER_URL = ('http://127.0.0.1:%d/ServeConstrained?' %
                 cns._DEFAULT_SERVING_PORT)

  # Setting for latency testing.
  _LATENCY_TEST_SECS = 1

  def _StartServer(self):
    """Starts the CNS, returns pid."""
    cmd = ['python', 'cns.py', '--interface=%s' % _INTERFACE]
    process = subprocess.Popen(cmd, stderr=subprocess.PIPE)

    # Wait for server to startup.
    line = True
    while line:
      line = process.stderr.readline()
      if 'STARTED' in line:
        return process.pid

    self.fail('Failed to start CNS.')

  def setUp(self):
    # Start the CNS.
    self._server_pid = self._StartServer()

    # Create temp file for serving. Run after server start so if a failure
    # during setUp() occurs we don't leave junk files around.
    f, self._file = tempfile.mkstemp(dir=os.getcwd())
    os.write(f, self._TEST_DATA)
    os.close(f)

    # Strip cwd off so we have a proper relative path.
    self._relative_fn = self._file[len(os.getcwd()) + 1:]

  def tearDown(self):
    os.unlink(self._file)
    os.kill(self._server_pid, signal.SIGTERM)

  def testServerServesFiles(self):
    now = time.time()

    f = urllib2.urlopen('%sf=%s' % (self._SERVER_URL, self._relative_fn))

    # Verify file data is served correctly.
    self.assertEqual(self._TEST_DATA, f.read())

    # For completeness ensure an unconstrained call takes less time than our
    # artificial constraints checked in the tests below.
    self.assertTrue(time.time() - now < self._LATENCY_TEST_SECS)

  def testServerLatencyConstraint(self):
    """Tests serving a file with a latency network constraint."""
    # Abort if does not have root access.
    self.assertEqual(os.geteuid(), 0, 'You need root access to run this test.')
    now = time.time()

    base_url = '%sf=%s' % (self._SERVER_URL, self._relative_fn)
    url = '%s&latency=%d' % (base_url, self._LATENCY_TEST_SECS * 1000)
    f = urllib2.urlopen(url)

    # Verify file data is served correctly.
    self.assertEqual(self._TEST_DATA, f.read())

    # Verify the request took longer than the requested latency.
    self.assertTrue(time.time() - now > self._LATENCY_TEST_SECS)

    # Verify the server properly redirected the URL.
    self.assertTrue(f.geturl().startswith(base_url.replace(
        str(cns._DEFAULT_SERVING_PORT), str(cns._DEFAULT_CNS_PORT_RANGE[0]))))


class ConstrainedNetworkServerUnitTests(unittest.TestCase):
  """ConstrainedNetworkServer class unit tests."""

  def testGetServerURL(self):
    """Test server URL is correct when using Cherrypy port."""
    cns_obj = cns.ConstrainedNetworkServer(self.DummyOptions(), None)

    self.assertEqual(cns_obj._GetServerURL('ab/xz.webm', port=1234, t=1),
                     'http://127.0.0.1:1234/ServeConstrained?f=ab/xz.webm&t=1')

  def testGetServerURLWithLocalServer(self):
    """Test server URL is correct when using --local-server-port port."""
    cns_obj = cns.ConstrainedNetworkServer(self.DummyOptionsWithServer(), None)

    self.assertEqual(cns_obj._GetServerURL('ab/xz.webm', port=1234, t=1),
                     'http://127.0.0.1:1234/media/ab/xz.webm?t=1')

  class DummyOptions(object):
    www_root = 'media'
    port = 9000
    cherrypy.url = lambda: 'http://127.0.0.1:9000/ServeConstrained'
    local_server_port = None

  class DummyOptionsWithServer(object):
    www_root = 'media'
    port = 9000
    cherrypy.url = lambda: 'http://127.0.0.1:9000/ServeConstrained'
    local_server_port = 8080


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