From ec99ad22b5957ab6af3f0938864b088ac8c75d9d Mon Sep 17 00:00:00 2001
From: Cayo Puigdefabregas <cayo@ribaguifi.com>
Date: Thu, 17 Mar 2022 11:50:57 +0100
Subject: [PATCH] .

---
 ereuse_devicehub/client.py | 60 ++++++++++++++++++++++++++++++++++++++
 tests/conftest.py          | 27 ++++++++++++++++-
 tests/test_render_2_0.py   | 34 ++++++++++++---------
 3 files changed, 107 insertions(+), 14 deletions(-)

diff --git a/ereuse_devicehub/client.py b/ereuse_devicehub/client.py
index 3224dfd3..3ddf94ff 100644
--- a/ereuse_devicehub/client.py
+++ b/ereuse_devicehub/client.py
@@ -1,5 +1,7 @@
 from inspect import isclass
 from typing import Dict, Iterable, Type, Union
+from flask.testing import FlaskClient
+from flask_wtf.csrf import generate_csrf
 
 from ereuse_utils.test import JSON, Res
 from teal.client import Client as TealClient, Query, Status
@@ -156,3 +158,61 @@ class UserClient(Client):
         response = super().login(self.email, self.password)
         self.user = response[0]
         return response
+
+
+class UserClientFlask:
+
+    def __init__(self, application,
+                 email: str,
+                 password: str,
+                 response_wrapper=None,
+                 use_cookies=True,
+                 follow_redirects=True):
+        self.email = email
+        self.password = password
+        self.follow_redirects = follow_redirects
+        self.user = None
+
+        self.client = FlaskClient(application, use_cookies=use_cookies)
+        self.client.get('/login/')
+
+        data = {
+            'email': email,
+            'password': password,
+            'csrf_token': generate_csrf(),
+        }
+        body, status, headers = self.client.post('/login/', data=data, follow_redirects=True)
+        self.headers = headers
+        body = next(body).decode("utf-8")
+        assert "Unassgined" in body
+
+    def get(self,
+            uri='',
+            data=None,
+            follow_redirects=True,
+            **kw):
+
+        body, status, headers = self.client.get(
+            uri,
+            data=data,
+            follow_redirects=follow_redirects,
+            headers=self.headers
+        )
+        body = next(body).decode("utf-8")
+        return (body, status)
+
+    def post(self,
+            uri='',
+            data=None,
+            follow_redirects=True,
+            **kw):
+
+        import pdb; pdb.set_trace()
+        body, status, headers = self.client.post(
+            uri,
+            data=data,
+            follow_redirects=follow_redirects,
+            headers=self.headers
+        )
+        body = next(body).decode("utf-8")
+        return (body, status)
diff --git a/tests/conftest.py b/tests/conftest.py
index b85b34fd..e571e96b 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -13,7 +13,7 @@ from decouple import config
 from psycopg2 import IntegrityError
 from sqlalchemy.exc import ProgrammingError
 
-from ereuse_devicehub.client import Client, UserClient
+from ereuse_devicehub.client import Client, UserClient, UserClientFlask
 from ereuse_devicehub.config import DevicehubConfig
 from ereuse_devicehub.db import db
 from ereuse_devicehub.devicehub import Devicehub
@@ -136,6 +136,31 @@ def user2(app: Devicehub) -> UserClient:
         return client
 
 
+@pytest.fixture()
+def user3(app: Devicehub) -> UserClientFlask:
+    """Gets a client with a logged-in dummy user."""
+    with app.app_context():
+        password = 'foo'
+        user = create_user(password=password)
+        client = UserClientFlask(
+            app, user.email, password
+        )
+        return client
+
+
+@pytest.fixture()
+def user4(app: Devicehub) -> UserClient:
+    """Gets a client with a logged-in dummy user."""
+    with app.app_context():
+        password = 'foo'
+        email = 'foo2@foo.com'
+        user = create_user(email=email, password=password)
+        client = UserClientFlask(
+            app, user.email, password
+        )
+        return client
+
+
 def create_user(email='foo@foo.com', password='foo') -> User:
     user = User(email=email, password=password)
     user.individuals.add(Person(name='Timmy'))
diff --git a/tests/test_render_2_0.py b/tests/test_render_2_0.py
index 0a3a19ad..3d3775ed 100644
--- a/tests/test_render_2_0.py
+++ b/tests/test_render_2_0.py
@@ -2,7 +2,7 @@ import pytest
 from flask.testing import FlaskClient
 from flask_wtf.csrf import generate_csrf
 
-from ereuse_devicehub.client import UserClient
+from ereuse_devicehub.client import UserClient, UserClientFlask
 from ereuse_devicehub.devicehub import Devicehub
 from tests import conftest
 
@@ -35,19 +35,27 @@ def test_login(user: UserClient, app: Devicehub):
 
 @pytest.mark.mvp
 @pytest.mark.usefixtures(conftest.app_context.__name__)
-def test_inventory(user: UserClient, app: Devicehub):
-    client = FlaskClient(app, use_cookies=True)
-    client.get('/login/')
+def test_inventory(user3: UserClientFlask):
+    body, status = user3.get('/inventory/device/')
+
+    assert status == '200 OK'
+    assert "Unassgined" in body
+
+
+@pytest.mark.mvp
+@pytest.mark.usefixtures(conftest.app_context.__name__)
+def test_add_lot(user3: UserClientFlask):
+    body, status = user3.get('/inventory/lot/add/')
+
+    assert status == '200 OK'
+    assert "Add a new lot" in body
+
     data = {
-        'email': user.email,
-        'password': 'foo',
-        'remember': False,
+        'name': 'lot1',
         'csrf_token': generate_csrf(),
     }
-    body, status, headers = client.post('/login/', data=data, follow_redirects=True)
-    body, status, headers = client.get('/inventory/device/', headers=headers)
-
-    body = next(body).decode("utf-8")
-    assert status == '200 OK'
     # import pdb; pdb.set_trace()
-    assert "Unassgined" in body
+    body, status = user3.post('/inventory/lot/add/', data=data)
+
+    assert status == '200 OK'
+    assert "lot1" in body