File diff 897541fd8978 → a69e9814cf16
conntrackt/tests/factories.py
Show inline comments
 
new file 100644
 
# Third-party application imports.
 
import factory
 

	
 
# Application imports.
 
from conntrackt.models import Project, Location, Entity, Interface, Communication
 

	
 

	
 
class ProjectFactory(factory.django.DjangoModelFactory):
 
    """
 
    Factory for producing projects where name is set to "Test Project N", and N
 
    is an increasing sequence. Description is based on name.
 
    """
 

	
 
    FACTORY_FOR = Project
 

	
 
    @factory.sequence
 
    def name(n):
 
        return "Test Project %d" % n
 

	
 
    description = factory.LazyAttribute(lambda d: "This is %s." % d.name)
 

	
 

	
 
class LocationFactory(factory.django.DjangoModelFactory):
 
    """
 
    Factory for producing locations where name is set to "Test Location N", and
 
    N is an increasing sequence. Description is based on name.
 
    """
 

	
 
    FACTORY_FOR = Location
 

	
 
    @factory.sequence
 
    def name(n):
 
        return "Test Location %d" % n
 

	
 
    description = factory.LazyAttribute(lambda d: "This is %s." % d.name)
 

	
 

	
 
class ServerInterfaceFactory(factory.django.DjangoModelFactory):
 
    """
 
    Factory for producing server interfaces. Interface IP address is set to
 
    192.168.1.N, where N is a sequence. The netmask is always set to
 
    "255.255.255.255". The interface name is always set to eth0.
 
    """
 

	
 
    FACTORY_FOR = Interface
 
    netmask = "255.255.255.255"
 
    name = "eth0"
 

	
 
    @factory.sequence
 
    def address(n):
 
        return "192.168.1.%d" % (n)
 

	
 

	
 
class SubnetInterfaceFactory(factory.django.DjangoModelFactory):
 
    """
 
    Factory for producing subnet "interfaces". Interface IP address is set to
 
    10.10.N.0, where N is a sequence. Interface name is always set to
 
    switch0. Netmask is set to "255.255.255.0".
 
    """
 

	
 
    FACTORY_FOR = Interface
 
    netmask = "255.255.255.0"
 
    name = "switch0"
 

	
 
    @factory.sequence
 
    def address(n):
 
        return "10.10.%d.0" % (n)
 

	
 

	
 
class ServerEntityFactory(factory.django.DjangoModelFactory):
 
    """
 
    Factory for producing server entities where name is set to "Test Entity N",
 
    and N is an increasing sequence.
 
    """
 

	
 
    FACTORY_FOR = Entity
 

	
 
    interface = factory.RelatedFactory(ServerInterfaceFactory, "entity")
 

	
 
    @factory.sequence
 
    def name(n):
 
        return "Test Entity %d" % n
 

	
 
    description = factory.LazyAttribute(lambda d: "This is %s." % d.name)
 

	
 

	
 
class SubnetEntityFactory(factory.django.DjangoModelFactory):
 
    """
 
    Factory for producing subnet entities where name is set to "Test Subnet N",
 
    and N is an increasing sequence.
 
    """
 

	
 
    FACTORY_FOR = Entity
 

	
 
    interface = factory.RelatedFactory(SubnetInterfaceFactory, "entity")
 

	
 
    @factory.sequence
 
    def name(n):
 
        return "Test Subnet %d" % n
 

	
 
    description = factory.LazyAttribute(lambda d: "This is %s." % d.name)
 

	
 

	
 
class CommunicationFactory(factory.django.DjangoModelFactory):
 
    """
 
    Factory for producing communications. The descriptin of communication will
 
    be based on protocol and port.
 
    """
 

	
 
    FACTORY_FOR = Communication
 

	
 
    description = factory.LazyAttribute(lambda d: "Communicate over %s:%s" % (d.protocol, d.port))
 

	
 

	
 
def setup_test_data():
 
    """
 
    Sets-up some test data for testing more complex functionality.
 
    """
 

	
 
    for factory in ProjectFactory, LocationFactory, ServerEntityFactory, SubnetEntityFactory, ServerInterfaceFactory, SubnetInterfaceFactory, CommunicationFactory:
 
        factory.reset_sequence()
 

	
 
    project1 = ProjectFactory(pk=1)
 
    project2 = ProjectFactory(pk=2)
 

	
 
    location1 = LocationFactory(pk=1)
 
    location2 = LocationFactory(pk=2)
 

	
 
    entity1 = ServerEntityFactory(pk=1, project=project1, location=location1)
 
    entity2 = ServerEntityFactory(pk=2, project=project1, location=location1)
 
    entity3 = ServerEntityFactory(pk=3, project=project1, location=location2)
 
    entity4 = SubnetEntityFactory(pk=4, project=project1, location=location2)
 
    entity5 = ServerEntityFactory(pk=5, project=project2, location=location1)
 

	
 
    communication1 = CommunicationFactory(pk=1, source_id=2, destination_id=1, protocol="TCP", port="22")
 
    communication2 = CommunicationFactory(pk=2, source_id=2, destination_id=1, protocol="ICMP", port="8")
 
    communication3 = CommunicationFactory(pk=3, source_id=3, destination_id=1, protocol="TCP", port="3306")
 
    communication4 = CommunicationFactory(pk=4, source_id=1, destination_id=3, protocol="UDP", port="53")
 
    communication5 = CommunicationFactory(pk=5, source_id=4, destination_id=1, protocol="TCP", port="22")
 
    communication6 = CommunicationFactory(pk=6, source_id=1, destination_id=2, protocol="UDP", port="123")