def test_registration(self):
class TestMixin:
pass
self.assertTrue(Mixin.register('mixin', TestMixin))
self.assertTrue(Mixin.unregister('mixin'))
self.assertFalse(Mixin.unregister('mixin'))
def test_underunder_enter_and_underunder_exit(self):
class A(object):
pass
class B(A):
pass
for mixin in self.expected_mixins:
if mixin:
mixee = Mixin(A(), mixin)
mixee.__enter__()
self.assertTrue(mixin in A.__bases__)
self.assertFalse(mixin in B.__bases__)
mixee.__exit__()
self.assertFalse(mixin in A.__bases__)
self.assertFalse(mixin in B.__bases__)
def test_get_mixin(self):
keys = 'hi!', 'i have', 'heard of?', None
actual_mixins = [Mixin.get_mixin(key) for key in keys]
for ex, ac in zip(self.expected_mixins, actual_mixins):
self.assertEqual(ex, ac)
self.default_validator()
assert payload['type']
assert payload['type'] in valid_payload_types
except AssertionError:
raise RequestError("Invalid payload: %s" % payload)
return fn(self, **kwargs)
return validator
class Query:
""" Query handler """
@validate
def execute(self, host=None):
""" Handle an inbound query """
key, value = self.packet.payload['type'], self.packet.payload['value']
result = []
for host in self._r.keys(pattern="host_*"):
host = json.JSONDecoder().decode(self._r.get(host))
query_values = host.get(key, None)
if query_values:
if value in query_values:
result.append(host)
return result
Mixin.register("heard of?", Query)
""" Validates a packet"""
payload = self.packet.payload
try:
valid_payload_types = ('hosts', 'services', 'addresses', 'disks')
self.default_validator()
print payload
assert payload['type']
assert payload['type'] in valid_payload_types
assert type(payload['export']) == list
assert len(payload['export']) > 0
except AssertionError:
raise RequestError("Invalid payload: %s" % payload)
return fn(self, **kwargs)
return validator
class Export:
""" Export handler """
@validate
def execute(self, host=None):
""" handle an inbound export """
if not host:
host = self.packet.payload['host']
key = 'host_%s' % host
return self._store(key, self.packet.payload)
Mixin.register("i have", Export)
from pynpoint.errors import RequestError
def validate(fn):
""" Validation decorator """
def validator(self, **kwargs):
""" Validates a packet"""
payload = self.packet.payload
try:
self.default_validator()
assert payload['host'] or kwargs['host']
except AssertionError:
raise RequestError("Invalid payload: %s" % payload)
return fn(self, **kwargs)
return validator
class Announcement:
""" Announcement RequestHandler mixin """
@validate
def execute(self, host=None):
""" handle a hello """
if not host:
host = self.packet.payload['host']
key = 'host_%s' % host
return self._store(key, self.packet.payload)
Mixin.register("hi!", Announcement)
def handle(self):
mixin = Mixin.get_mixin(self.packet.request_type)
with Mixin(self, mixin) as handler:
# handler.execute is really self, with mixin mixed in.
return handler.execute()
