diff --git a/dataserv/Contract.py b/dataserv/Contract.py index a98e883..b057032 100644 --- 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 index 28bbcac..abeccee 100644 --- 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 index bdc695a..bb27594 100644 --- 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 index da0a447..73b6389 100644 --- 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) \ No newline at end of file