Working storage function

rival_regions_wrapper/__init__.py

@@ -8,4 +8,4 @@ of some Rival Regions functionalities.
 
 from .authentication_handler import AuthenticationHandler
 from .middleware import LocalAuthentication, RemoteAuthentication
-from .api_wrapper import Profile
+from .api_wrapper import Profile, Storage

rival_regions_wrapper/api_wrapper/__init__.py

@@ -30,3 +30,4 @@ if None in (USERNAME, PASSWORD, LOGIN_METHOD):
 MIDDLEWARE = LocalAuthentication(USERNAME,PASSWORD,LOGIN_METHOD)
 
 from .profile import Profile
+from .storage import Storage

rival_regions_wrapper/api_wrapper/profile.py

@@ -19,7 +19,6 @@ class Profile(object):
         soup = BeautifulSoup(response, 'html.parser')
         level = soup.select_one('div.oil:nth-child(2) > div:nth-child(2)').text
         perks = soup.select('table tr:nth-child(2) span')
-        print(perks)
         profile = {
             'profile_id': self.profile_id,
             'name': re.sub(r'.*:\s', '', soup.find('h1').text),

rival_regions_wrapper/api_wrapper/storage.py

@@ -0,0 +1,46 @@
+"""Profile class"""
+
+import re
+
+from bs4 import BeautifulSoup
+
+from . import MIDDLEWARE
+
+
+class Storage(object):
+    """Wrapper class for profile"""
+
+    @staticmethod
+    def info():
+        """Get profile"""
+        path = 'storage'
+        response = MIDDLEWARE.get(path)
+        soup = BeautifulSoup(response, 'html.parser')
+        keys = {
+            'oil': 3,
+            'ore': 4,
+            'uranium': 11,
+            'diamonds': 15,
+            'liquid_oxygen': 21,
+            'helium-3': 24,
+            'rivalium': 26,
+            'antirad': 13,
+            'energy_drink': 17,
+            'spacerockets': 20,
+            'lss': 25,
+            'tanks': 2,
+            'aircrafts': 1,
+            'missiles': 14,
+            'bombers': 16,
+            'battleships': 18,
+            'laser_drones': 27,
+            'moon_tanks': 22,
+            'space_stations': 23
+        }
+        storage = {}
+        for key, item_id in keys.items():
+            storage[key] = int(soup.find('span', {'urlbar' : item_id}).text.replace('.', ''))
+            storage['{}_max'.format(key)] = int(soup.find('span', {'urlbar' : item_id})['maxstore'])
+
+        print(storage)
+        return storage

tests/test_rival_regions_wrapper.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from rival_regions_wrapper.api_wrapper import Profile
+from rival_regions_wrapper.api_wrapper import Profile, Storage
 
 
 @pytest.fixture
@@ -13,10 +13,8 @@ def profile_keys():
 @pytest.mark.vcr()
 def test_profile_info(profile_keys):
     """Test an API call to get client info"""
-
     profile_instance = Profile(192852686)
     response = profile_instance.info()
-    print(response)
 
     assert isinstance(response, dict), "The response should be a dict"
     assert response['profile_id'] == 192852686, "The ID should be in the response"
@@ -32,7 +30,20 @@ def test_profile_info(profile_keys):
 def storage_keys():
     """Standard keys for storage"""
     return [
-        'profile_id', 'oil', 'ore', 'uranium', 'diamonds', 'liquid_oxygen', 'rivalium',
-        'antirad', 'energy_drink', 'spacerockets', 'lss', 'tanks', 'aircrafts',
-        'missiles', 'bombers', 'battleships', 'laser_drones', 'moon_tanks', 'space_stations'
+        'oil', 'ore', 'uranium', 'diamonds', 'liquid_oxygen',
+        'helium-3', 'rivalium', 'antirad', 'energy_drink', 
+        'spacerockets', 'lss', 'tanks', 'aircrafts', 'missiles', 
+        'bombers', 'battleships', 'laser_drones', 'moon_tanks', 'space_stations',
+        'oil_max', 'ore_max', 'uranium_max', 'diamonds_max', 'liquid_oxygen_max',
+        'helium-3_max', 'rivalium_max', 'antirad_max', 'energy_drink_max', 
+        'spacerockets_max', 'lss_max', 'tanks_max', 'aircrafts_max', 'missiles_max', 
+        'bombers_max', 'battleships_max', 'laser_drones_max', 'moon_tanks_max', 'space_stations'
     ]
+
+@pytest.mark.vcr()
+def test_storage_info(storage_keys):
+    """Test an API call to get storage info"""
+    response = Storage.info()
+
+    assert isinstance(response, dict), "The response should be a dict"
+    assert set(storage_keys).issubset(response.keys()), "All keys should be in the response"