Add retires to context initializers

3 files changed, 13 insertions, 6 deletions

diff --git a/tftpy/TftpContexts.py b/tftpy/TftpContexts.py
index da85886..679413e 100644
--- a/tftpy/TftpContexts.py
+++ b/tftpy/TftpContexts.py
@@ -76,7 +76,7 @@ class TftpMetrics(object):
 class TftpContext(object):
     """The base class of the contexts."""
 
-    def __init__(self, host, port, timeout, localip = ""):
+    def __init__(self, host, port, timeout, retries=DEF_TIMEOUT_RETRIES, localip = ""):
         """Constructor for the base context, setting shared instance
         variables."""
         self.file_to_transfer = None
@@ -88,6 +88,7 @@ class TftpContext(object):
             self.sock.bind((localip, 0))
         self.sock.settimeout(timeout)
         self.timeout = timeout
+        self.retries = retries
         self.state = None
         self.next_block = 0
         self.factory = TftpPacketFactory()
@@ -212,11 +213,13 @@ class TftpContextServer(TftpContext):
                  timeout,
                  root,
                  dyn_file_func=None,
-                 upload_open=None):
+                 upload_open=None,
+                 retries=DEF_TIMEOUT_RETRIES):
         TftpContext.__init__(self,
                              host,
                              port,
                              timeout,
+                             retries
                              )
         # At this point we have no idea if this is a download or an upload. We
         # need to let the start state determine that.
@@ -267,11 +270,13 @@ class TftpContextClientUpload(TftpContext):
                  options,
                  packethook,
                  timeout,
+                 retries=DEF_TIMEOUT_RETRIES,
                  localip = ""):
         TftpContext.__init__(self,
                              host,
                              port,
                              timeout,
+                             retries,
                              localip)
         self.file_to_transfer = filename
         self.options = options
@@ -320,7 +325,7 @@ class TftpContextClientUpload(TftpContext):
             except TftpTimeout as err:
                 log.error(str(err))
                 self.retry_count += 1
-                if self.retry_count >= TIMEOUT_RETRIES:
+                if self.retry_count >= self.retries:
                     log.debug("hit max retries, giving up")
                     raise
                 else:
@@ -346,11 +351,13 @@ class TftpContextClientDownload(TftpContext):
                  options,
                  packethook,
                  timeout,
+                 retries=DEF_TIMEOUT_RETRIES,
                  localip = ""):
         TftpContext.__init__(self,
                              host,
                              port,
                              timeout,
+                             retries,
                              localip)
         # FIXME: should we refactor setting of these params?
         self.file_to_transfer = filename
@@ -403,7 +410,7 @@ class TftpContextClientDownload(TftpContext):
             except TftpTimeout as err:
                 log.error(str(err))
                 self.retry_count += 1
-                if self.retry_count >= TIMEOUT_RETRIES:
+                if self.retry_count >= self.retries:
                     log.debug("hit max retries, giving up")
                     raise
                 else:
diff --git a/tftpy/TftpServer.py b/tftpy/TftpServer.py
index 95ca70e..f121d21 100644
--- a/tftpy/TftpServer.py
+++ b/tftpy/TftpServer.py
@@ -210,7 +210,7 @@ class TftpServer(TftpSession):
                 except TftpTimeout as err:
                     log.error(str(err))
                     self.sessions[key].retry_count += 1
-                    if self.sessions[key].retry_count >= TIMEOUT_RETRIES:
+                    if self.sessions[key].retry_count >= self.sessions[key].retries:
                         log.debug("hit max retries on %s, giving up" %
                             self.sessions[key])
                         deletion_list.append(key)
diff --git a/tftpy/TftpShared.py b/tftpy/TftpShared.py
index 88530c3..734b263 100644
--- a/tftpy/TftpShared.py
+++ b/tftpy/TftpShared.py
@@ -9,7 +9,7 @@ DEF_BLKSIZE = 512
 MAX_BLKSIZE = 65536
 SOCK_TIMEOUT = 5
 MAX_DUPS = 20
-TIMEOUT_RETRIES = 5
+DEF_TIMEOUT_RETRIES = 5
 DEF_TFTP_PORT = 69
 
 # A hook for deliberately introducing delay in testing.