From 22646ca76258e215655ff15a12d7974d8bd1840a Mon Sep 17 00:00:00 2001
From: Mario Rial <mrial@datto.com>
Date: Thu, 16 May 2019 02:58:41 -0400
Subject: [PATCH] Test improvements

---
 channels/testing/http.py   | 35 +++++++++++++++++++++++++++++-
 tests/test_generic_http.py | 44 ++++++++++++++++++++++----------------
 2 files changed, 59 insertions(+), 20 deletions(-)

diff --git a/channels/testing/http.py b/channels/testing/http.py
index 6b1514ca7..760d522ca 100644
--- a/channels/testing/http.py
+++ b/channels/testing/http.py
@@ -31,7 +31,7 @@ def __init__(self, application, method, path, body=b"", headers=None):
 
     async def get_response(self, timeout=1):
         """
-        Get the application's response. Returns a dict with keys of
+        Get the application's full response. Returns a dict with keys of
         "body", "headers" and "status".
         """
         # If we've not sent the request yet, do so
@@ -54,3 +54,36 @@ async def get_response(self, timeout=1):
         del response_start["type"]
         response_start.setdefault("headers", [])
         return response_start
+
+    async def send_request(self):
+        """
+        Sends the request to the application without then waiting for
+         headers or any response.
+        """
+        if not self.sent_request:
+            self.sent_request = True
+            await self.send_input({"type": "http.request", "body": self.body})
+
+    async def get_response_start(self, timeout=1):
+        """
+        Gets the start of the response (its headers and status code)
+        """
+        response_start = await self.receive_output(timeout)
+        assert response_start["type"] == "http.response.start"
+
+        # Return structured info
+        del response_start["type"]
+        response_start.setdefault("headers", [])
+        return response_start
+
+    async def get_body_chunk(self, timeout=1):
+        """
+        Gets one chunk of body.
+        """
+        chunk = await self.receive_output(timeout)
+        assert chunk["type"] == "http.response.body"
+        assert isinstance(chunk["body"], bytes)
+        if not chunk.get("more_body", False):
+            await self.wait(timeout)
+
+        return chunk["body"]
diff --git a/tests/test_generic_http.py b/tests/test_generic_http.py
index 80920b4bd..2d773360e 100644
--- a/tests/test_generic_http.py
+++ b/tests/test_generic_http.py
@@ -1,4 +1,3 @@
-import asyncio
 import json
 
 import pytest
@@ -48,11 +47,11 @@ class TestConsumer(AsyncHttpConsumer):
         Abstract consumer that provides a method that handles running a command and getting a response on a
         device.
         """
-
-        groups = ['test_group']
+        channel_layer_alias = 'testlayer'
 
         async def handle(self, body):
-            print("Latest Channel ID: ")
+            # Add consumer to a known test group that we will use to send events to.
+            await self.channel_layer.group_add('test_group', self.channel_name)
             await self.send_headers(
                 status=200,
                 headers=[
@@ -61,13 +60,14 @@ async def handle(self, body):
             )
 
         async def send_to_long_poll(self, event):
-            print("RunCommandConsumer: Event received on send to long poll.")
-            command_output = str(event['data']).encode('utf8')
-            await self.send_body(command_output, more_body=False)
+            received_data = str(event['data']).encode('utf8')
+            # We just echo what we receive, and close the response.
+            await self.send_body(received_data, more_body=False)
 
     channel_layers_setting = {
-        "default": {"BACKEND": "channels.layers.InMemoryChannelLayer"}
+        'testlayer': {"BACKEND": "channels.layers.InMemoryChannelLayer"}
     }
+
     with override_settings(CHANNEL_LAYERS=channel_layers_setting):
         # Open a connection
         communicator = HttpCommunicator(
@@ -77,14 +77,20 @@ async def send_to_long_poll(self, event):
             body=json.dumps({"value": 42, "anything": False}).encode("utf-8"),
         )
 
-        def send_to_channel_layer():
-            print("send to channel layer called")
-            channel_layer = get_channel_layer()
-            # Test that the websocket channel was added to the group on connect
-            message = {"type": "send.to.long.poll", "data": "hello"}
-            asyncio.ensure_future(channel_layer.group_send("chat", message))
-
-        asyncio.get_event_loop().call_later(3, send_to_channel_layer)
-        print("Making http requests.")
-        response = await communicator.get_response(timeout=10)
-        assert True
+        # We issue the HTTP request
+        await communicator.send_request()
+
+        # Gets the response start (status and headers)
+        response_start = await communicator.get_response_start(timeout=1)
+
+        # Make sure that the start of the response looks good so far.
+        assert response_start['status'] == 200
+        assert response_start['headers'] == [(b'Content-Type', b'application/json')]
+
+        # Send now a message to the consumer through the channel layer. Using the known test_group.
+        channel_layer = get_channel_layer('testlayer')
+        await channel_layer.group_send('test_group', {"type": "send.to.long.poll", "data": "hello from channel layer"})
+
+        # Now we should be able to get the message back on the remaining chunk of body.
+        body = await communicator.get_body_chunk(timeout=1)
+        assert body == b"hello from channel layer"