Started adding tests for WFQ scheduler

Author: Axel Dahlberg

qlinklayer/distQueue.py

@@ -3,15 +3,15 @@
 #
 # Implements a simple distributed queue shared with one other node over a connection.
 #
-# Author: Stephanie Wehner
+# Authors: Stephanie Wehner, Matthew Skrzypczyk, Axel Dahlberg
 
 from copy import copy
 from collections import deque, defaultdict
 from functools import partial
 from easysquid.easyfibre import ClassicalFibreConnection
 from easysquid.easyprotocol import EasyProtocol, ClassicalProtocol
 from netsquid.pydynaa import EventType, EventHandler
-from qlinklayer.localQueue import TimeoutLocalQueue, EGPLocalQueue
+from qlinklayer.localQueue import TimeoutLocalQueue, EGPLocalQueue, WFQLocalQueue
 from qlinklayer.toolbox import LinkLayerException
 from easysquid.toolbox import logger
 
@@ -1094,3 +1094,22 @@ def _post_process_cmd_ADD_ACK(self, qid, qseq):
             Sequence number within local queue of item
         """
         self.queueList[qid].ack(qseq)
+
+
+class WFQDistributedQueue(EGPDistributedQueue):
+    def _init_queues(self, numQueues=1, maxSeq=2 ** 8, throw_local_queue_events=False):
+        """
+        Initializes the local queues
+        :param numQueues: int
+            Number of queues
+        :param throw_local_queue_events: bool
+            Whether the local queues should throw events for data collection or not
+        :return:
+        """
+        # Initialize queues
+        self.queueList = []
+        self.numLocalQueues = numQueues
+        for j in range(numQueues):
+            q = WFQLocalQueue(qid=j, maxSeq=maxSeq, throw_events=throw_local_queue_events,
+                              timeout_callback=self.timeout_callback)
+            self.queueList.append(q)

qlinklayer/localQueue.py

@@ -306,6 +306,19 @@ def get_new_queue_item(self, request, seq):
         return _EGPLocalQueueItem(request, seq, self.qid, self.timeout_callback)
 
 
+class WFQLocalQueue(EGPLocalQueue):
+    def get_new_queue_item(self, request, seq):
+        """
+        Returns a fresh queue item instance
+        :param request: :obj:`qlinklayer.EGP.EGPRequest`
+            The request
+        :param seq: int
+            The queue item sequence number
+        :return: :obj:`qlinklayer.localQueue._EGPLocalQueueItem`
+        """
+        return _WFQLocalQueueItem(request, seq, self.qid, self.timeout_callback)
+
+
 class TimeoutLocalQueue(LocalQueue):
     def __init__(self, qid=None, wsize=None, maxSeq=None, throw_events=False):
         """

tests/test_scheduler.py

@@ -1,16 +1,19 @@
 import unittest
 import logging
+from easysquid.easyfibre import ClassicalFibreConnection
 from easysquid.easynetwork import EasyNetwork
 from easysquid.qnode import QuantumNode
 from easysquid.quantumMemoryDevice import NVCommunicationDevice
 from easysquid.easyprotocol import TimedProtocol
 from easysquid.toolbox import logger
 from netsquid.simutil import sim_run, sim_reset
-from qlinklayer.distQueue import EGPDistributedQueue
-from qlinklayer.scheduler import StrictPriorityRequestScheduler
+from qlinklayer.distQueue import EGPDistributedQueue, WFQDistributedQueue
+from qlinklayer.scheduler import StrictPriorityRequestScheduler, WFQRequestScheduler
 from qlinklayer.qmm import QuantumMemoryManagement
 from qlinklayer.egp import EGPRequest
 from qlinklayer.scenario import EGPSimulationScenario
+from qlinklayer.feu import SingleClickFidelityEstimationUnit
+from qlinklayer.mhp import SimulatedNodeCentricMHPService
 
 logger.setLevel(logging.CRITICAL)
 
@@ -150,6 +153,27 @@ def test_priority(self):
             test_scheduler.clear_request(next_aid)
 
 
+class TestWFQRequestScheduler(unittest.TestCase):
+    def setUp(self):
+        alice = QuantumNode("alice", 0)
+        bob = QuantumNode("bob", 0)
+        conn = ClassicalFibreConnection(alice, bob, length=0.001)
+
+        distQueueA = WFQDistributedQueue(alice, numQueues=3, accept_all=True)
+        distQueueB = WFQDistributedQueue(bob, numQueues=3, accept_all=True)
+        distQueueA.connect_to_peer_protocol(distQueueB, conn=conn)
+
+        qmmA = QuantumMemoryManagement(alice)
+        qmmB = QuantumMemoryManagement(bob)
+
+        # mhp_service = SimulatedNodeCentricMHPService("mhp_service", alice, bob)
+        #
+        # feuA = SingleClickFidelityEstimationUnit(alice, )
+        # self.schedulerA = WFQRequestScheduler(distQueueA, qmmA, feuA)
+    def test_init(self):
+        pass
+
+
 class TestTimings(unittest.TestCase):
     def setUp(self):
         def timeout_handler(request):