"""
<Program Name>
  z_test_client.repy

<Started>
  February 13, 2009

<Purpose>
  To test that Client functions work together.

  Passes quietly
"""
include client.repy

if callfunc == 'initialize':
  DEST_IP = 'whatever'
  DEST_PORT = 12345
  SRC_IP = 'whateversrc'
  SRC_PORT = 12346
  TIMEOUT = .1
  RETRIES = 4  
  MSS = 2
  CONG_WINDOW_SIZE = MSS
  HUGE = MSS * 100

def test_client():
  mycontext['packet_count'] = 0
  mycontext['packets'].clear()
  mycontext['packetc'].clear()
  
  SEQ_NUM = 2
  ACK_NUM = SEQ_NUM + 1
  DATA = "h"
  SND_WINDOW= 10
  RCV_WINDOW= 20

  mycontext['client'] = TcpClient(ack_all_register)
  cli = mycontext['client']
  cli.last_seq_sent = SEQ_NUM
  cli.last_ack_recv = ACK_NUM
  cli.recv_window_size= RCV_WINDOW
  cli.send_window_size= SND_WINDOW
  bytes = cli.send(DATA, DEST_PORT, SRC_PORT,RETRIES, TIMEOUT)

  assert bytes == len(DATA)
  assert  len(mycontext['packets']) == 1
  assert ACK_NUM in mycontext['packets'].keys()
  assert mycontext['packets'][ACK_NUM].is_payload()
  assert cli.last_ack_recv == ACK_NUM + len(DATA)
  assert cli.last_seq_sent == SEQ_NUM + len(DATA)
  assert mycontext['packet_count'] == 1
  assert mycontext['packets'][ACK_NUM].srcport == SRC_PORT
  assert mycontext['packets'][ACK_NUM].destport == DEST_PORT
  assert mycontext['packets'][ACK_NUM].window == RCV_WINDOW
  assert mycontext['packets'][ACK_NUM].ack_num == 0
  assert mycontext['packets'][ACK_NUM].seq_num == ACK_NUM


def test_client_window():
  mycontext['packet_count'] = 0
  mycontext['packets'].clear()
  mycontext['packetc'].clear()
  
  SEQ_NUM = 2
  ACK_NUM = SEQ_NUM + 1
  DATA = "hihihihi"
  SND_WINDOW= HUGE
  RCV_WINDOW= HUGE
  CONG_WIN = HUGE

  mycontext['client'] = TcpClient(ack_all_register)
  cli = mycontext['client']
  cli.last_seq_sent = SEQ_NUM
  cli.last_ack_recv = ACK_NUM
  cli.recv_window_size= RCV_WINDOW
  cli.send_window_size= SND_WINDOW
  cli.cong_controller.window_size= CONG_WIN
  cli.send(DATA,DEST_PORT, SRC_PORT,RETRIES, TIMEOUT)

  assert  len(mycontext['packets']) == len(DATA)/MSS
  assert ACK_NUM in mycontext['packets'].keys()
  assert (ACK_NUM + MSS) in mycontext['packets'].keys()
  assert (ACK_NUM + 2 * MSS) in mycontext['packets'].keys()
  assert (ACK_NUM + 3 * MSS) in mycontext['packets'].keys()

  assert cli.last_ack_recv == ACK_NUM + len(DATA)
  assert cli.last_seq_sent == SEQ_NUM + len(DATA)

def test_client_sliding_window():
  mycontext['packet_count'] = 0
  mycontext['packets'].clear()
  mycontext['packetc'].clear()
  
  SEQ_NUM = 2
  ACK_NUM = SEQ_NUM + 1
  DATA = "hihihihi"
  SND_WINDOW= HUGE
  RCV_WINDOW= HUGE
  CONG_WIN = HUGE

  mycontext['client'] = TcpClient(ack_first_register)
  cli = mycontext['client']
  cli.last_seq_sent = SEQ_NUM
  cli.last_ack_recv = ACK_NUM
  cli.recv_window_size= RCV_WINDOW
  cli.send_window_size= SND_WINDOW
  cli.cong_controller.window_size= CONG_WIN
  cli.send(DATA, DEST_PORT, SRC_PORT,RETRIES, TIMEOUT)

  assert  len(mycontext['packets']) == len(DATA)/MSS
  assert ACK_NUM in mycontext['packetc'].keys()
  assert (ACK_NUM + MSS) in mycontext['packetc'].keys()
  assert (ACK_NUM + 2 * MSS) in mycontext['packetc'].keys()
  assert (ACK_NUM + 3 * MSS) in mycontext['packetc'].keys()

  assert mycontext['packetc'][ACK_NUM] == 1
  assert mycontext['packetc'][ACK_NUM + 1 * MSS] == 2
  assert mycontext['packetc'][ACK_NUM + 2 * MSS] == 2
  assert mycontext['packetc'][ACK_NUM + 3 * MSS] == 2

  assert cli.last_ack_recv == ACK_NUM + len(DATA)
  assert cli.last_seq_sent == SEQ_NUM + len(DATA)

def test_client_congestion_window_slow_start():
  SEQ_NUM = 2
  ACK_NUM = SEQ_NUM + 1
  DATA = "hihihihihihihihi" # eight packets
  SND_WINDOW= HUGE
  RCV_WINDOW= HUGE
  CONG_WIN = CONG_WINDOW_SIZE

  mycontext['client'] = TcpClient(ack_all_register)
  cli = mycontext['client']
  cli.last_seq_sent = SEQ_NUM
  cli.last_ack_recv = ACK_NUM
  cli.recv_window_size= RCV_WINDOW
  cli.send_window_size= SND_WINDOW
  cli.cong_controller.window_size= CONG_WIN
  cli.send(DATA, DEST_PORT, SRC_PORT, RETRIES, TIMEOUT)

  # Pk  1 2 3 4 5 6 7 8
  # CW  1 2 2 4 4 4 4 8
  assert cli.cong_controller.window_size == (len(DATA)/MSS + 1)*CONG_WINDOW_SIZE
  assert cli.cong_controller.threshold == None

def test_client_congestion_window_avoidance():
  SEQ_NUM = 2
  ACK_NUM = SEQ_NUM + 1
  DATA = "hihihihihihihihihi" # 9 packets
  SND_WINDOW= HUGE
  RCV_WINDOW= HUGE
  CONG_WIN = CONG_WINDOW_SIZE

  mycontext['client'] = TcpClient(ack_all_register)
  cli = mycontext['client']
  cli.last_seq_sent = SEQ_NUM
  cli.last_ack_recv = ACK_NUM
  cli.recv_window_size = RCV_WINDOW
  cli.send_window_size = SND_WINDOW
  cli.cong_controller.window_size = CONG_WIN * 2
  cli.cong_controller.threshold = CONG_WIN * 2
  cli.send(DATA, DEST_PORT, SRC_PORT, RETRIES, TIMEOUT)

  assert cli.cong_controller.threshold ==  CONG_WIN * 2
  assert cli.cong_controller.window_size == 5 * CONG_WINDOW_SIZE

def test_client_congestion_window_threshold():
  SEQ_NUM = 2
  ACK_NUM = SEQ_NUM + 1
  DATA = "hihi" # 2 packets
  SND_WINDOW= HUGE
  RCV_WINDOW= HUGE
  CONG_WIN = CONG_WINDOW_SIZE

  mycontext['client'] = TcpClient(ack_first_register)
  cli = mycontext['client']
  cli.last_seq_sent = SEQ_NUM
  cli.last_ack_recv = ACK_NUM
  cli.recv_window_size= RCV_WINDOW
  cli.send_window_size= SND_WINDOW
  cli.cong_controller.window_size= CONG_WIN * 3
  cli.cong_controller.threshold = None
  cli.send(DATA, DEST_PORT, SRC_PORT, RETRIES, TIMEOUT)

  assert cli.cong_controller.threshold ==  CONG_WIN * 2
  assert cli.cong_controller.window_size == CONG_WIN * 2

def test_client_congestion_window_threshold_reset():
  SEQ_NUM = 2
  ACK_NUM = SEQ_NUM + 1
  DATA = "hihi" # 2 packets
  SND_WINDOW= HUGE
  RCV_WINDOW= HUGE
  CONG_WIN = CONG_WINDOW_SIZE

  mycontext['client'] = TcpClient(ack_first_register)
  cli = mycontext['client']
  cli.last_seq_sent = SEQ_NUM
  cli.last_ack_recv = ACK_NUM
  cli.recv_window_size= RCV_WINDOW
  cli.send_window_size= SND_WINDOW
  cli.cong_controller.window_size= CONG_WIN * 3
  cli.cong_controller.threshold = CONG_WIN
  cli.send(DATA, DEST_PORT, SRC_PORT, RETRIES, TIMEOUT)

  assert cli.cong_controller.threshold ==  int(((CONG_WIN * 3.0) + (CONG_WIN/3.0))/2.0)
  assert cli.cong_controller.window_size == CONG_WIN * 2


####################
# Test Helper Functions
####################
def ack_all_register(message):
    pack = unpack(message)
    mycontext['packet_count'] +=1
    if not pack.seq_num in mycontext['packetc'].keys():
      mycontext['packetc'][pack.seq_num] = 1
    else:
      mycontext['packetc'][pack.seq_num] += 1
    mycontext['packets'][pack.seq_num] = pack
    bytes = len(pack.payload)

    # generate ACK yes-man style
    ack_pack = pack.copy()
    ack_pack.srcport =   pack.destport
    ack_pack.destport = pack.srcport
    ack_pack.seq_num = 0
    ack_pack.ack_num = pack.seq_num + bytes
    ack_pack.window = HUGE
    ack_pack.control_bits = "ACK"
    ack_pack.payload = "echo ACK-ing"
    settimer(0, mycontext['client'].process, (ack_pack,))

    return bytes

def ack_first_register(message):
    pack = unpack(message)
    mycontext['packet_count'] +=1
    if not pack.seq_num in mycontext['packetc'].keys():
      mycontext['packetc'][pack.seq_num] = 1
    else:
      mycontext['packetc'][pack.seq_num] += 1
    mycontext['packets'][pack.seq_num] = pack
    bytes = len(pack.payload)

    # the first packet in window
    if mycontext['client'].last_ack_recv == pack.seq_num:
      ack_pack = pack.copy()
      ack_pack.srcport =   pack.destport
      ack_pack.destport = pack.srcport
      ack_pack.seq_num = 0
      ack_pack.ack_num = pack.seq_num + bytes
      ack_pack.window = HUGE
      ack_pack.control_bits = "ACK"
      ack_pack.payload = "echo ACK-ing"
      settimer(0, mycontext['client'].process, (ack_pack,))
    else: # drop it  
      pass

    return bytes

def main():
  # setup
  mycontext['packet_count'] = 0
  mycontext['packets'] = {}
  mycontext['packetc'] = {} # packet counter 
  mycontext['client'] = None
  
  # Run tests
  test_client()
  test_client_window()
  test_client_sliding_window()
  test_client_congestion_window_slow_start()
  test_client_congestion_window_avoidance()
  test_client_congestion_window_threshold()
  test_client_congestion_window_threshold_reset()

  # teardown
  exitall()

if callfunc == 'initialize':
  main()
