From ebf1538313d682d291cfbcee3d1b822209a8006a Mon Sep 17 00:00:00 2001
From: Dan Helfman <dan@luminotes.com>
Date: Thu, 13 Sep 2007 00:53:49 +0000
Subject: [PATCH] model.User now has a storage_bytes member, and
 controller.Users has a method for calculating and updating this field.

---
 controller/Users.py           | 10 ++++++++++
 controller/test/Test_users.py | 11 +++++++++++
 model/User.py                 | 12 ++++++++++++
 model/test/Test_user.py       |  9 +++++++++
 4 files changed, 42 insertions(+)

diff --git a/controller/Users.py b/controller/Users.py
index 9916a2a..5dceff2 100644
--- a/controller/Users.py
+++ b/controller/Users.py
@@ -343,4 +343,14 @@ class Users( object ):
 
     return total_bytes
 
+  @async
+  def update_storage( self, user ):
+    """
+    Calculate and record total storage utilization for the given user.
+    @type user: User
+    @param user: user for which to calculate storage utilization
+    """
+    user.storage_bytes = self.calculate_storage( user )
+    yield None
+
   scheduler = property( lambda self: self.__scheduler )
diff --git a/controller/test/Test_users.py b/controller/test/Test_users.py
index 58843dd..e93b5fd 100644
--- a/controller/test/Test_users.py
+++ b/controller/test/Test_users.py
@@ -225,3 +225,14 @@ class Test_users( Test_controller ):
       self.database.size( self.anon_notebook.notes[ 0 ].object_id, self.anon_notebook.notes[ 0 ].revision )
 
     assert size == expected_size
+
+  def test_update_storage( self ):
+    previous_revision = self.user.revision
+
+    cherrypy.root.users.update_storage( self.user )
+    self.scheduler.wait_until_idle()
+
+    expected_size = cherrypy.root.users.calculate_storage( self.user )
+
+    assert self.user.storage_bytes == expected_size
+    assert self.user.revision > previous_revision
diff --git a/model/User.py b/model/User.py
index 83a0b8b..a00b77e 100644
--- a/model/User.py
+++ b/model/User.py
@@ -11,6 +11,12 @@ class User( Persistent ):
   SALT_CHARS = [ chr( c ) for c in range( ord( "!" ), ord( "~" ) + 1 ) ]
   SALT_SIZE = 12
 
+  def __setstate__( self, state ):
+    if "_User__storage_bytes" not in state:
+      state[ "_User__storage_bytes" ] = False
+
+    self.__dict__.update( state )
+
   def __init__( self, id, username, password, email_address, notebooks = None ):
     """
     Create a new user with the given credentials and information.
@@ -33,6 +39,7 @@ class User( Persistent ):
     self.__password_hash = self.__hash_password( password )
     self.__email_address = email_address
     self.__notebooks = notebooks or []
+    self.__storage_bytes = 0 # total storage bytes for this user's notebooks, notes, and revisions
 
   def __create_salt( self ):
     return "".join( [ random.choice( self.SALT_CHARS ) for i in range( self.SALT_SIZE ) ] )
@@ -88,9 +95,14 @@ class User( Persistent ):
     self.update_revision()
     self.__notebooks = notebooks
 
+  def __set_storage_bytes( self, storage_bytes ):
+    self.update_revision()
+    self.__storage_bytes = storage_bytes
+
   username = property( lambda self: self.secondary_id )
   email_address = property( lambda self: self.__email_address )
   password = property( None, __set_password )
+  storage_bytes = property( lambda self: self.__storage_bytes, __set_storage_bytes )
 
   # the notebooks (read-only and read-write) that this user has access to
   notebooks = property( lambda self: copy( self.__notebooks ), __set_notebooks )
diff --git a/model/test/Test_user.py b/model/test/Test_user.py
index 6f5fa20..532ff25 100644
--- a/model/test/Test_user.py
+++ b/model/test/Test_user.py
@@ -16,6 +16,7 @@ class Test_user( object ):
     assert self.user.username == self.username
     assert self.user.email_address == self.email_address
     assert self.user.notebooks == []
+    assert self.user.storage_bytes == 0
 
   def test_check_correct_password( self ):
     assert self.user.check_password( self.password ) == True
@@ -51,6 +52,14 @@ class Test_user( object ):
     assert self.user.notebooks[ 0 ].object_id == notebook_id
     assert self.user.revision > previous_revision
 
+  def test_set_storage_bytes( self ):
+    previous_revision = self.user.revision
+    storage_bytes = 44
+    self.user.storage_bytes = storage_bytes
+    
+    assert self.user.storage_bytes == storage_bytes
+    assert self.user.revision > previous_revision
+
 
 class Test_user_with_notebooks( object ):
   def setUp( self ):