1 files changed, 17 insertions, 8 deletions

diff --git a/tests/sockets/test_sockets_echo_client.c b/tests/sockets/test_sockets_echo_client.c
index f6ea85cf..18cff97e 100644
--- a/tests/sockets/test_sockets_echo_client.c
+++ b/tests/sockets/test_sockets_echo_client.c
@@ -71,21 +71,25 @@ void main_loop(void *arg) {
       return;
     }
 
+#if !TEST_DGRAM
     // as a test, confirm with ioctl that we have data available
     // after selecting
     int available;
     res = ioctl(server.fd, FIONREAD, &available);
     assert(res != -1);
     assert(available);
+#endif
 
     res = do_msg_read(server.fd, &server.msg, echo_read, 0, NULL, NULL);
-    if (res == 0) {
+    if (res == -1) {
+      return;
+    } else if (res == 0) {
       perror("server closed");
       finish(EXIT_FAILURE);
-    } else if (res != -1) {
-      echo_read += res;
     }
 
+    echo_read += res;
+
     // once we've read the entire message, validate it
     if (echo_read >= server.msg.length) {
       assert(!strcmp(server.msg.buffer, MESSAGE));
@@ -97,13 +101,15 @@ void main_loop(void *arg) {
     }
 
     res = do_msg_write(server.fd, &echo_msg, echo_wrote, 0, NULL, 0);
-    if (res == 0) {
+    if (res == -1) {
+      return;
+    } else if (res == 0) {
       perror("server closed");
       finish(EXIT_FAILURE);
-    } else if (res != -1) {
-      echo_wrote += res;
     }
 
+    echo_wrote += res;
+
     // once we're done writing the message, read it back
     if (echo_wrote >= echo_msg.length) {
       server.state = MSG_READ;
@@ -124,8 +130,11 @@ int main() {
   echo_msg.buffer = malloc(echo_msg.length);
   strncpy(echo_msg.buffer, MESSAGE, echo_msg.length);
 
+  echo_read = 0;
+  echo_wrote = 0;
+
   // create the socket and set to non-blocking
-#if !USE_UDP
+#if !TEST_DGRAM
   server.fd = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
 #else
   server.fd = socket(PF_INET, SOCK_DGRAM, IPPROTO_UDP);
@@ -158,4 +167,4 @@ int main() {
 #endif
 
   return EXIT_SUCCESS;
-}
\ No newline at end of file
+}