From 22b4079d501c39925c40772238fc942ac84fbc01 Mon Sep 17 00:00:00 2001
From: Rohith Asrk <rohith.asrk@gmail.com>
Date: Tue, 31 Jul 2018 15:19:11 +0530
Subject: [PATCH] Removed files related to batch user creation on deletion and
 changed filepaths

Closes #156 and #157
---
 django_freeradius/base/models.py            | 12 ++++++++++++
 django_freeradius/tests/base/test_api.py    |  5 +++--
 django_freeradius/tests/base/test_models.py | 12 +++++-------
 django_freeradius/utils.py                  |  3 ++-
 4 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/django_freeradius/base/models.py b/django_freeradius/base/models.py
index c8840f3..d18ba51 100644
--- a/django_freeradius/base/models.py
+++ b/django_freeradius/base/models.py
@@ -1,4 +1,5 @@
 import csv
+import os
 from base64 import encodestring
 from hashlib import md5, sha1
 from io import StringIO
@@ -617,6 +618,7 @@ def prefix_add(self, prefix, n, password_length=BATCH_DEFAULT_PASSWORD_LENGTH):
             u.save()
             self.users.add(u)
         f = generate_pdf(prefix, {'users': user_password})
+        f.name = f.name.split('/')[-1]
         self.pdf = f
         self.full_clean()
         self.save()
@@ -624,6 +626,7 @@ def prefix_add(self, prefix, n, password_length=BATCH_DEFAULT_PASSWORD_LENGTH):
     def delete(self):
         self.users.all().delete()
         super(AbstractRadiusBatch, self).delete()
+        self._remove_files()
 
     def expire(self):
         users = self.users.all()
@@ -631,6 +634,15 @@ def expire(self):
             u.is_active = False
             u.save()
 
+    def _remove_files(self):
+        strategy_filemap = {
+            'prefix': 'pdf',
+            'csv': 'csvfile'
+        }
+        path = getattr(self, strategy_filemap.get(self.strategy)).path
+        if os.path.isfile(path):
+            os.remove(path)
+
 
 @python_2_unicode_compatible
 class AbstractRadiusProfile(TimeStampedEditableModel):
diff --git a/django_freeradius/tests/base/test_api.py b/django_freeradius/tests/base/test_api.py
index 4489098..32bc2be 100644
--- a/django_freeradius/tests/base/test_api.py
+++ b/django_freeradius/tests/base/test_api.py
@@ -1,5 +1,6 @@
 import json
 
+from django.conf import settings
 from django.contrib.auth.models import User
 from django.urls import reverse
 from django.utils.timezone import now
@@ -560,10 +561,10 @@ def test_batch_csv_201(self):
         self.assertEqual(self.radius_batch_model.objects.count(), 0)
         self.assertEqual(User.objects.count(), 0)
         text = 'user,cleartext$abcd,email@gmail.com,firstname,lastname'
-        with open('test.csv', 'wb') as file:
+        with open('{}/test.csv'.format(settings.MEDIA_ROOT), 'wb') as file:
             text2 = text.encode('utf-8')
             file.write(text2)
-        with open('test.csv', 'rb') as file:
+        with open('{}/test.csv'.format(settings.MEDIA_ROOT), 'rb') as file:
             data = {
                 "name": "test",
                 "strategy": "csv",
diff --git a/django_freeradius/tests/base/test_models.py b/django_freeradius/tests/base/test_models.py
index 39e7dee..a938aa4 100644
--- a/django_freeradius/tests/base/test_models.py
+++ b/django_freeradius/tests/base/test_models.py
@@ -61,14 +61,12 @@ def test_string_representation(self):
         radiusbatch = self.radius_batch_model(name='test')
         self.assertEqual(str(radiusbatch), 'test')
 
-    def test_custom_queryset(self):
-        radiusbatch = self.radius_batch_model.objects.create()
+    def test_delete_method(self):
+        radiusbatch = self.radius_batch_model.objects.create(strategy='prefix',
+                                                             prefix='test',
+                                                             name='test')
+        radiusbatch.prefix_add('test', 5)
         User = get_user_model()
-        for i in range(5):
-            user = User.objects.create(username='rohith{}'.format(str(i)))
-            user.set_password('password')
-            user.save()
-            radiusbatch.users.add(user)
         self.assertEqual(User.objects.all().count(), 5)
         radiusbatch.delete()
         self.assertEqual(self.radius_batch_model.objects.all().count(), 0)
diff --git a/django_freeradius/utils.py b/django_freeradius/utils.py
index bdc8f7b..6e66bd9 100644
--- a/django_freeradius/utils.py
+++ b/django_freeradius/utils.py
@@ -2,6 +2,7 @@
 from io import StringIO
 
 import swapper
+from django.conf import settings
 from django.contrib.auth import get_user_model
 from django.core.exceptions import ValidationError
 from django.core.files import File
@@ -63,7 +64,7 @@ def prefix_generate_users(prefix, n, password_length):
 def generate_pdf(prefix, data):
     template = get_template(BATCH_PDF_TEMPLATE)
     html = template.render(data)
-    f = open('{}.pdf'.format(prefix), 'w+b')
+    f = open('{}/{}.pdf'.format(settings.MEDIA_ROOT, prefix), 'w+b')
     pisa.CreatePDF(html.encode('utf-8'), dest=f, encoding='utf-8')
     f.seek(0)
     return File(f)