diff --git a/app/models/user.rb b/app/models/user.rb
index f385a69c..84085994 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -225,6 +225,14 @@ def brief
     end
   end
 
+  def team
+    if team_id
+      Team.find(team_id)
+    else
+      membership.try(:team)
+    end
+  end
+
   def team_ids
     [team_id]
   end
@@ -248,7 +256,7 @@ def teams_being_followed
   end
 
   def on_team?
-    not team_document_id.nil?
+    team_id.present? || membership.present?
   end
 
   def team_member_of?(user)
diff --git a/spec/models/user_spec.rb b/spec/models/user_spec.rb
index f39d87d3..9b9e63ac 100644
--- a/spec/models/user_spec.rb
+++ b/spec/models/user_spec.rb
@@ -214,6 +214,23 @@ class AlsoNotaBadge < BadgeBase
     end
   end
 
+  describe '#team' do
+    let(:team) { Fabricate(:team) }
+    let(:user) { Fabricate(:user) }
+    it 'returns membership team if user has membership' do
+      team.add_member(user)
+      expect(user.team).to eq(team)
+    end
+    it 'returns team if team_id is set' do
+      user.team_id = team.id
+      user.save
+      expect(user.team).to eq(team)
+    end
+    it 'returns nil if no team_id or membership' do
+      expect(user.team).to eq(nil)
+    end
+  end
+
   it 'should indicate when user is on a premium team' do
     team = Fabricate(:team, premium: true)
     member = team.add_member(user = Fabricate(:user))
@@ -234,6 +251,22 @@ class AlsoNotaBadge < BadgeBase
     expect(user.team).to be_nil
   end
 
+  describe '#on_team?' do
+    let(:team) { Fabricate(:team) }
+    let(:user) { Fabricate(:user) }
+    it 'is true if user has a membership' do
+      expect(user.on_team?).to eq(false)
+      team.add_member(user)
+      expect(user.reload.on_team?).to eq(true)
+    end
+    it 'is true if user is on a team' do
+      expect(user.on_team?).to eq(false)
+      user.team = team
+      user.save
+      expect(user.reload.on_team?).to eq(true)
+    end
+  end
+
   it 'can follow another user' do
     user = Fabricate(:user)
     other_user = Fabricate(:user)