more testing

5 files changed, 39 insertions, 11 deletions

diff --git a/.travis.yml b/.travis.yml
index 2e83771..a91ecba 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -9,8 +9,9 @@ branches:
 dist: xenial
 sudo: required
 
-#services:
-#  - docker
+services:
+  - rabbitmq
+  #- docker
 
 before_install:
   - sudo apt-get update
@@ -23,6 +24,7 @@ env:
     - PIPENV_IGNORE_VIRTUALENVS=1
     - PIPENV_NOSPIN=1
     - PIPENV_HIDE_EMOJIS=1
+    - RABBITMQ_HOST=localhost
 
 install:
   - pip install pipenv
diff --git a/README.md b/README.md
index 59e86ae..a58223c 100644
--- a/README.md
+++ b/README.md
@@ -39,6 +39,7 @@ result <- |             |<----------|            |<--------|           | <------
 install python packages
 
 ```bash
+apt-get install -y libprotobuf-dev #needed by nsjail
 pipenv sync --dev
 ```
 
diff --git a/config.py b/config.py
index beb435e..684412e 100644
--- a/config.py
+++ b/config.py
@@ -12,7 +12,7 @@ def autodiscover():
                 host = list(container.attrs.get('NetworkSettings').get('Networks').values())[0]['IPAddress']
                 return host
     except Exception:
-        return '172.17.0.2'
+        return '127.0.0.1'
 
 
 USERNAME = os.environ.get('RMQ_USERNAME', 'rabbits')
diff --git a/rmq.py b/rmq.py
index c93e787..ef5e8c2 100644
--- a/rmq.py
+++ b/rmq.py
@@ -37,12 +37,11 @@ class Rmq(object):
     def _declare(self, channel, queue):
         channel.queue_declare(
             queue=queue,
-            durable=False,                      # Do not commit messages to disk
+            durable=False,                       # Do not commit messages to disk
             arguments={'x-message-ttl': 5000},  # Delete message automatically after x milliseconds
-            auto_delete=True)                   # Delete queue when all connection are closed
-
-    def consume(self, queue=QUEUE, callback=None, thread_ws=None):
+            auto_delete=True)                    # Delete queue when all connection are closed
 
+    def consume(self, queue=QUEUE, callback=None, thread_ws=None, run_once=False):
         while True:
             try:
                 connection = pika.BlockingConnection(self.con_params)
@@ -51,16 +50,22 @@ class Rmq(object):
                     channel = connection.channel()
                     self._declare(channel, queue)
                     channel.basic_qos(prefetch_count=1)
-                    channel.basic_consume(
-                        lambda ch, method, properties, body:
-                        callback(ch, method, properties, body, thread_ws=thread_ws),
-                        queue=queue)
+
+                    if not run_once:
+                        channel.basic_consume(
+                            lambda ch, method, properties, body:
+                            callback(ch, method, properties, body, thread_ws=thread_ws),
+                            queue=queue)
 
                     log.info(f"Connected to host: {self.host} port: {self.port} queue: {queue}")
 
                     if thread_ws:
                         if not thread_ws.closed:
                             thread_ws.send('{"service": "connected"}')
+
+                    if run_once:
+                        return channel.basic_get(queue=queue)
+
                     channel.start_consuming()
 
                 except Exception:
@@ -114,6 +119,9 @@ class Rmq(object):
                     log.info((f"published: {self.host} "
                               f"queue: {queue} "
                               f"message: {message}"))
+
+                    return result
+
                 else:
                     log.error(f"Message '{message}' not delivered")
 
diff --git a/tests/test_snekbox.py b/tests/test_snekbox.py
index 4d86ecf..d2f81fe 100644
--- a/tests/test_snekbox.py
+++ b/tests/test_snekbox.py
@@ -1,12 +1,18 @@
 import unittest
 import pytest
 import os
+import json
 
 from snekbox import Snekbox
+from rmq import Rmq
+
+r = Rmq()
+
 python_binary = os.environ.get('PYTHONEXECUTABLE', '/usr/bin/python3.6')
 nsjail = os.sep.join([os.getcwd(), f'binaries{os.sep}nsjail2.6-ubuntu-x86_64'])
 snek = Snekbox(nsjail_binary=nsjail, python_binary=python_binary)
 
+
 class SnekTests(unittest.TestCase):
 	def test_nsjail(self):
 		result = snek.python3('print("test")')
@@ -39,3 +45,14 @@ class SnekTests(unittest.TestCase):
 			self.assertIn('ModuleNotFoundError', result.strip())
 		else:
 			self.assertIn('returned non-zero exit status 1.', result.strip())
+
+
+class RMQTests(unittest.TestCase):
+	def test_a_publish(self):
+		message = json.dumps({"snekid": "test", "message": "print('test')"})
+		result = r.publish(message)
+		self.assertTrue(result)
+
+	def test_b_consume(self):
+		result = r.consume(callback=snek.message_handler, queue='input', run_once=True)
+		self.assertEquals(result[2], b'{"snekid": "test", "message": "print(\'test\')"}')