3 files changed, 27 insertions, 5 deletions
diff --git a/test/unit/tmp_user_test.rb b/test/unit/tmp_user_test.rb
index af2218e..0a9ad68 100644
--- a/test/unit/tmp_user_test.rb
+++ b/test/unit/tmp_user_test.rb
@@ -2,19 +2,23 @@ require 'test_helper'
 
 class TmpUserTest < ActiveSupport::TestCase
 
+  setup do
+    InviteCodeValidator.any_instance.stubs(:not_existent?).returns(false)
+  end
+
   test "test_user saved to tmp_users" do
     begin
       assert User.ancestors.include?(TemporaryUser)
 
       assert_difference('User.database.info["doc_count"]') do
         normal_user = User.create!(:login => 'a'+SecureRandom.hex(5).downcase,
-          :password_verifier => 'ABCDEF0010101', :password_salt => 'ABCDEF', :invite_code => 'testcode')
+          :password_verifier => 'ABCDEF0010101', :password_salt => 'ABCDEF')
         refute normal_user.database.to_s.include?('tmp')
       end
 
       assert_difference('User.tmp_database.info["doc_count"]') do
         tmp_user = User.create!(:login => 'test_user_'+SecureRandom.hex(5).downcase,
-          :password_verifier => 'ABCDEF0010101', :password_salt => 'ABCDEF', :invite_code => 'testcode')
+          :password_verifier => 'ABCDEF0010101', :password_salt => 'ABCDEF')
         assert tmp_user.database.to_s.include?('tmp')
       end
     ensure
diff --git a/test/unit/token_test.rb b/test/unit/token_test.rb
index 5468650..dcf3fc4 100644
--- a/test/unit/token_test.rb
+++ b/test/unit/token_test.rb
@@ -4,6 +4,7 @@ class TokenTest < ActiveSupport::TestCase
   include StubRecordHelper
 
   setup do
+    InviteCodeValidator.any_instance.stubs(:not_existent?).returns(false)
     @user = find_record :user
   end
 
diff --git a/test/unit/user_test.rb b/test/unit/user_test.rb
index cd290d5..e28bef6 100644
--- a/test/unit/user_test.rb
+++ b/test/unit/user_test.rb
@@ -4,6 +4,7 @@ class UserTest < ActiveSupport::TestCase
 
   include SRP::Util
   setup do
+    InviteCodeValidator.any_instance.stubs(:not_existent?).returns(false)
     @user = FactoryGirl.build(:user)
   end
 
@@ -70,9 +71,25 @@ class UserTest < ActiveSupport::TestCase
     assert_equal key, @user.public_key
   end
 
-  test "user should have an invite token" do
-    user = User.new
-    assert_nil(user.invite_code)
+  test "user should not be created with invalid invite code" do
+    InviteCodeValidator.any_instance.stubs(:not_existent?).returns(true)
+    invalid_user = FactoryGirl.build(:user)
+
+    assert !invalid_user.valid?
+  end
+
+  test "user should be created with valid invite code" do
+    assert @user.valid?
+  end
+
+  test "trying to create a user with invalid invite code should add error" do
+    InviteCodeValidator.any_instance.stubs(:not_existent?).returns(true)
+    invalid_user = FactoryGirl.build(:user)
+
+    invalid_user.valid?
+
+    errors = {invite_code: ["This is not a valid code"]}
+    assert_equal errors, invalid_user.errors.messages
   end
 
   #