added set_bandwidth tests

StorjOld · Mar 10, 2016 · 14cb3b0 · 14cb3b0
1 parent a89bc6b
commit 14cb3b0
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 6 deletions.
diff --git a/dataserv/app.py b/dataserv/app.py
@@ -167,13 +167,8 @@ def set_bandwidth(btc_addr, bandwidth):
     try:
         user = Farmer(btc_addr)
         user.authenticate(dict(request.headers))
-        assert(bandwidth >= 0)
         user.set_bandwidth(bandwidth, ip=request.remote_addr)
         return make_response("Bandwidth accepted.", 200)
-    except AssertionError:
-        msg = "Invalid bandwidth value: {0}".format(bandwidth)
-        logger.warning(msg)
-        return make_response(msg, 400)
     except ValueError:
         msg = "Invalid Bitcoin address."
         logger.warning(msg)

diff --git a/tests/test_App.py b/tests/test_App.py
@@ -231,6 +231,28 @@ def test_farmer_order(self):
         self.assertEqual(farmers[0].btc_addr, addr1)
 
 
+class BandwidthTest(TemplateTest):
+
+    def test_farmer_set_bandwidth(self):
+        # not found
+        btc_addr = self.gen_wallet()
+        rv = self.app.get('/api/bandwidth/{0}/1'.format(btc_addr))
+        self.assertEqual(rv.status_code, 404)
+
+        # register farmer
+        self.app.get('/api/register/{0}'.format(btc_addr))
+
+        # correct
+        rv = self.app.get('/api/bandwidth/{0}/5'.format(btc_addr))
+        self.assertEqual(rv.status_code, 200)
+        # rv = self.app.get('/api/online/json'.format(btc_addr))
+        # self.assertTrue(b"5" in rv.data)
+
+        # invalid btc address
+        rv = self.app.get('/api/bandwidth/{0}/1'.format(self.bad_addr))
+        self.assertEqual(rv.status_code, 400)
+
+
 class HeightTest(TemplateTest):
 
     def test_farmer_set_height(self):

diff --git a/tests/test_Farmer.py b/tests/test_Farmer.py
@@ -91,6 +91,21 @@ def test_ping_time_limit(self):
         delta_seconds = int((farmer.last_seen - register_time).seconds)
         self.assertEqual(delta_seconds, 0)
 
+    def test_bandwidth(self):
+        btc_addr = self.btctxstore.get_address(self.btctxstore.get_key(
+                                        self.btctxstore.create_wallet()))
+        farmer = Farmer(btc_addr)
+        farmer.register()
+
+        # set height and check function output
+        self.assertEqual(farmer.bandwidth, 0)
+        self.assertEqual(farmer.set_bandwidth(5), 5)
+        self.assertEqual(farmer.bandwidth, 5)
+
+        # check the db object as well
+        farmer2 = farmer.lookup()
+        self.assertEqual(farmer2.bandwidth, 5)
+
     def test_height(self):
         btc_addr = self.btctxstore.get_address(self.btctxstore.get_key(
                                         self.btctxstore.create_wallet()))
@@ -133,11 +148,12 @@ def test_to_json(self):
 
         farmer.ping()
         farmer.set_height(50)
+        farmer.set_bandwidth(55)
 
         test_json = {
             "height": 50,
             "ip": "",
-            "bandwidth": 0,
+            "bandwidth": 55,
             "btc_addr": btc_addr,
             'payout_addr': btc_addr,
             "last_seen": 0,