Travis Fix 3 + Extra Code

StorjOld · Jun 24, 2015 · beb3f08 · beb3f08
1 parent 43e2f51
commit beb3f08
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 4 deletions.
diff --git a/dataserv/Contract.py b/dataserv/Contract.py
@@ -43,6 +43,10 @@ def new_contract(self, seed=None, byte_size=None):
         """Build a new contract."""
         self.contract_type = 0
 
+        # make sure that farmer can actually create new contracts
+        if not self.below_limit():
+            raise MemoryError("Contract Capacity Limit Reached.")
+
         # take in a seed, if not generate it ourselves
         if seed is None:
             seed = os.urandom(12)

diff --git a/dataserv/app.py b/dataserv/app.py
@@ -1,7 +1,7 @@
 import sys
 import os.path
 import datetime
-from flask import make_response, jsonify
+from flask import Flask, make_response, jsonify
 sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
 
 
@@ -62,7 +62,7 @@ def ping(btc_addr):
         msg = "Invalid BTC Address."
         return make_response(error_msg.format(msg), 400)
     except LookupError:
-        msg = "Farmer not found."
+        msg = "Farmer Not Found."
         return make_response(error_msg.format(msg), 404)
 
 
@@ -107,8 +107,11 @@ def new_contract(btc_addr):
         msg = "Invalid BTC Address."
         return make_response(error_msg.format(msg), 400)
     except LookupError:
-        msg = "Farmer not found."
+        msg = "Farmer Not Found."
         return make_response(error_msg.format(msg), 404)
+    except MemoryError:
+        msg = "Contract Capacity Limit Reached."
+        return make_response(error_msg.format(msg), 413)
 
 
 if __name__ == '__main__':  # pragma: no cover

diff --git a/tests/test_App.py b/tests/test_App.py
@@ -64,7 +64,7 @@ def test_ping_not_found(self):
         rv = self.app.get('/api/ping/{0}'.format(addr))
 
         # good ping
-        self.assertEqual(b"Ping Failed: Farmer not found.", rv.data)
+        self.assertEqual(b"Ping Failed: Farmer Not Found.", rv.data)
         self.assertEqual(rv.status_code, 404)
 
     def test_ping_invalid_address(self):

diff --git a/tests/test_Contract.py b/tests/test_Contract.py
@@ -57,3 +57,8 @@ def test_contract_limit(self):
         con3.new_contract('b780bd4852b8c8a62859a50c', 1024)
         self.assertFalse(con3.below_limit(2000))
         self.assertTrue(con3.below_limit(3073))
+
+        # change config to break limit
+        app.config["BYTE_FARMER_MAX"] = 1024
+        with self.assertRaises(MemoryError):
+            con3.new_contract('b780bd4852b8c8a62859a50c', 1024)