From f512827589cda826b97f1f6cad3782374cf20ff4 Mon Sep 17 00:00:00 2001
From: dave <dave@davepedu.com>
Date: Mon, 12 Feb 2018 21:53:31 -0800
Subject: [PATCH] Better dcc tests

---
 pyircbot/modules/DCC.py   |  7 +++-
 tests/modules/test_dcc.py | 70 ++++++++++++++++++++++++++++++++++-----
 2 files changed, 68 insertions(+), 9 deletions(-)

diff --git a/pyircbot/modules/DCC.py b/pyircbot/modules/DCC.py
index 04b1086..ed972d0 100644
--- a/pyircbot/modules/DCC.py
+++ b/pyircbot/modules/DCC.py
@@ -17,6 +17,10 @@ from time import sleep
 BUFFSIZE = 8192
 
 
+class TransferFailedException(Exception):
+    pass
+
+
 def ip2int(ipstr):
     """
     Convert an ip address string to an integer
@@ -89,8 +93,9 @@ class RecieveGenerator(object):
                 yield chunk
                 if total >= self.length:
                     break
+            print("total", total, "expected", self.length)
             if total != self.length:
-                raise Exception("Transfer failed: expected {} bytes but got {}".format(self.length, total))
+                raise TransferFailedException("Transfer failed: expected {} bytes but got {}".format(self.length, total))
             raise StopIteration()
         finally:
             self.sock.shutdown(socket.SHUT_RDWR)
diff --git a/tests/modules/test_dcc.py b/tests/modules/test_dcc.py
index 0db7934..82ae2f8 100644
--- a/tests/modules/test_dcc.py
+++ b/tests/modules/test_dcc.py
@@ -2,8 +2,8 @@ import pytest
 import os
 import hashlib
 from time import sleep
-from pyircbot.modules.DCC import int2ip
 from tests.lib import *  # NOQA - fixtures
+from pyircbot.modules import DCC
 
 
 @pytest.fixture
@@ -20,26 +20,30 @@ def dccbot(fakebot):
     return fakebot
 
 
-def test_offerrecv(dccbot, tmpdir):
-    # allocate a temp file
-    flen = 1024 * 51
-    tmpfpath = os.path.join(tmpdir.dirname, "hello.bin")
+def _make_test_tempfiles(length, basedir):
+    tmpfpath = os.path.join(basedir.dirname, "hello.bin")
     with open(tmpfpath, 'wb') as fout:
-        fout.write(os.urandom(flen))
+        fout.write(os.urandom(length))
     # hash the tmpfile for later comparison
     m = hashlib.sha256()
     with open(tmpfpath, "rb") as ftmp:
         m.update(ftmp.read())
     srchash = m.hexdigest()
+    return tmpfpath, srchash
+
+
+def test_offerrecv(dccbot, tmpdir):
+    # allocate a temp file
+    flen = 1024 * 51
+    tmpfpath, srchash = _make_test_tempfiles(flen, tmpdir)
     # offer th file over DCC
     ip, port, reported_len, offer = dccbot.moduleInstances['DCC'].offer(tmpfpath)
     reported_len = int(reported_len)
     assert reported_len == flen, "offer reported wrong file length!"
-    ip = int2ip(ip)
+    ip = DCC.int2ip(ip)
     while not offer.bound:
         sleep(0.001)
     # receive the file over DCC
-    print("connecting to {}:{}".format(ip, port))
     recver = dccbot.moduleInstances['DCC'].recieve(ip, port, reported_len)
     data = b''
     d = hashlib.sha256()
@@ -49,3 +53,53 @@ def test_offerrecv(dccbot, tmpdir):
     # verify hashes and lengths
     assert len(data) == flen, "file not completely transferred"
     assert d.hexdigest() == srchash, "file was mangled in transfer"
+
+
+def test_tooshortfails(dccbot, tmpdir):
+        # allocate a temp file
+    flen = 1024 * 51
+    tmpfpath, srchash = _make_test_tempfiles(flen, tmpdir)
+    # offer th file over DCC
+    ip, port, reported_len, offer = dccbot.moduleInstances['DCC'].offer(tmpfpath)
+    reported_len = int(reported_len)
+    assert reported_len == flen, "offer reported wrong file length!"
+    ip = DCC.int2ip(ip)
+    while not offer.bound:
+        sleep(0.001)
+    recver = dccbot.moduleInstances['DCC'].recieve(ip, port, reported_len + 1)  # the magic
+    # fail to receive the file over DCC
+    # with pytest.raises(TransferFailedException):
+    try:
+        for chunk in iter(recver):
+            print(len(chunk))
+    # except DCC.TransferFailedException as te:
+    #    return
+    except Exception as fe:
+        assert fe.args == ('Transfer failed: expected 52225 bytes but got 52224',)
+        return
+    raise Exception("Did not raise")
+
+
+def test_toolongfails(dccbot, tmpdir):
+        # allocate a temp file
+    flen = 1024 * 51
+    tmpfpath, srchash = _make_test_tempfiles(flen, tmpdir)
+    # offer th file over DCC
+    ip, port, reported_len, offer = dccbot.moduleInstances['DCC'].offer(tmpfpath)
+    reported_len = int(reported_len)
+    assert reported_len == flen, "offer reported wrong file length!"
+    ip = DCC.int2ip(ip)
+    while not offer.bound:
+        sleep(0.001)
+    recver = dccbot.moduleInstances['DCC'].recieve(ip, port, reported_len - 1)  # the magic
+    # fail to receive the file over DCC
+    # with pytest.raises(TransferFailedException):
+    try:
+        for chunk in iter(recver):
+            print(len(chunk))
+    # except DCC.TransferFailedException as te:
+    #    return
+    except Exception as fe:
+        assert fe.args == ('Transfer failed: expected 52223 bytes but got 52224',)
+        return
+    raise Exception("Did not raise")