diff --git a/conntrackt/tests/factories.py b/conntrackt/tests/factories.py new file mode 100644 --- /dev/null +++ b/conntrackt/tests/factories.py @@ -0,0 +1,140 @@ +# 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")