def unitTest(cls): print "Testing QAgent" from system import System from networks import genGridNetwork from model import PJ s = System(genGridNetwork((4,3)),PJ()) s.cleanse() s.infections[2] = True s.infections[4] = True s.infections[7] = True numTrt = Agent.numTrt(s) numNodes = s.network.numNodes assert numNodes == 12, ("numNodes is not 12") assert numTrt == 3, ("numTrt is not 3") maxInfInd = 1 << numNodes maxTrtInd = Agent.numValidTrt(numNodes,numTrt) q = [[0]*maxTrtInd for i in range(maxInfInd)] for i in range(maxInfInd): for t in range(maxTrtInd): q[i][t] = t trtNodes = [9,10,11] a = cls.fromQ(q) a.applyTrt(s) if s.nTrt > len(trtNodes): raise ValueError("QAgent gave too many treatments ") elif s.nTrt < len(trtNodes): raise ValueError("QAgent gave too few treatments") else: for i in trtNodes: if not s.treatments[i]: raise ValueError("QAgent did not treat node " + i)