From a38f74e4e63601805a54ca68cb93f04589d72c61 Mon Sep 17 00:00:00 2001
From: alonbg <alon.blayergat@gmail.com>
Date: Mon, 22 Aug 2016 17:18:53 +0300
Subject: [PATCH] bugfix: Dgram: insert a receive block only if "---
 dgram_response" is specified.

---
 lib/Test/Nginx/Socket/Lua/Dgram.pm | 85 ++++++++++++++++++------------
 1 file changed, 50 insertions(+), 35 deletions(-)

diff --git a/lib/Test/Nginx/Socket/Lua/Dgram.pm b/lib/Test/Nginx/Socket/Lua/Dgram.pm
index 02a22d47..60a5337f 100644
--- a/lib/Test/Nginx/Socket/Lua/Dgram.pm
+++ b/lib/Test/Nginx/Socket/Lua/Dgram.pm
@@ -101,13 +101,11 @@ _EOC_
 _EOC_
         } else {
             $new_http_server_config .= <<_EOC_;
-                        sock:send('trigger dgram req')
-                        local data, err = sock:receive()
-                        if not data then
-                            sock:close()
-                            ngx.say("receive stream response error: ", err)
-                            return
-                        end
+                    local bytes, err = sock:send('trigger dgram req')
+                    if not bytes then
+                      ngx.say("send stream request error: ", err)
+                      return
+                    end
 _EOC_
         }
 
@@ -128,6 +126,14 @@ _EOC_
             || defined $block->dgram_response
             || defined $block->dgram_response_like)
         {
+            $new_http_server_config .= <<_EOC_;
+                    local data, err = sock:receive()
+                    if not data then
+                      	sock:close()
+                      	ngx.say("receive stream response error: ", err)
+                      	return
+                    end
+_EOC_
             if (defined $block->log_dgram_response) {
                 $new_http_server_config .= <<_EOC_;
                     print("stream response: ", data)
@@ -162,16 +168,15 @@ _EOC_
                         return
                     end
 _EOC_
-            }
-
-            $new_http_server_config .= <<_EOC_;
-                        sock:send('trigger_dgram_req2')
-                        local data, err = sock:receive()
-                        if not data then
-                            ngx.say("receive stream response error: ", err)
-                            return
-                        end
+            } else {
+                $new_http_server_config .= <<_EOC_;
+                    local bytes, err = sock:send('trigger dgram req 2')
+                    if not bytes then
+                      ngx.say("send stream request error: ", err)
+                      return
+                    end
 _EOC_
+            }
 
             if (defined $block->response_body
                 || defined $block->response_body_like
@@ -179,7 +184,13 @@ _EOC_
                 || defined $block->dgram_response_like)
             {
                 $new_http_server_config .= <<_EOC_;
-                        ngx.print(data)
+                    local data, err = sock:receive()
+                    if not data then
+                        sock:close()
+                        ngx.say("receive stream response error: ", err)
+                        return
+                    end
+                    ngx.print(data)
 _EOC_
             }
         }
@@ -197,24 +208,22 @@ _EOC_
             if (defined $dgram_req3) {
                 my $file = gen_data_file($dgram_req3);
                 $new_http_server_config .= <<_EOC_;
-                        local f = assert(io.open('$file', 'r'))
-                        local data = assert(f:read("*a"))
-                        assert(f:close())
-                        local bytes, err = sock:send(data)
-                        if not bytes then
-                            ngx.say("send stream request error: ", err)
-                            return
-                        end
+                    local f = assert(io.open('$file', 'r'))
+                    local data = assert(f:read("*a"))
+                    assert(f:close())
+                    local bytes, err = sock:send(data)
+                    if not bytes then
+                        ngx.say("send stream request error: ", err)
+                        return
+                    end
 _EOC_
-            }
-
-            $new_http_server_config .= <<_EOC_;
-                        sock:send('trigger_dgram_req3')
-                        local data, err = sock:receive()
-                        if not data then
-                            ngx.say("receive stream response error: ", err)
-                            return
-                        end
+            } else {
+                $new_http_server_config .= <<_EOC_;
+                    local bytes, err = sock:send('trigger dgram req 3')
+                    if not bytes then
+                      ngx.say("send stream request error: ", err)
+                      return
+                    end
 _EOC_
 
             if (defined $block->response_body
@@ -223,7 +232,13 @@ _EOC_
                 || defined $block->dgram_response_like)
             {
                 $new_http_server_config .= <<_EOC_;
-                        ngx.print(data)
+                    local data, err = sock:receive()
+                    if not data then
+                        sock:close()
+                        ngx.say("receive stream response error: ", err)
+                        return
+                    end
+                    ngx.print(data)
 _EOC_
             }
         }