diff --git a/pkg/gui/context.go b/pkg/gui/context.go
index 8d147f4c9..79f9acd94 100644
--- a/pkg/gui/context.go
+++ b/pkg/gui/context.go
@@ -43,13 +43,12 @@ func (gui *Gui) currentContextKeyIgnoringPopups() types.ContextKey {
 // use replaceContext when you don't want to return to the original context upon
 // hitting escape: you want to go that context's parent instead.
 func (gui *Gui) replaceContext(c types.Context) error {
-	gui.State.ContextManager.Lock()
-	defer gui.State.ContextManager.Unlock()
-
 	if !c.IsFocusable() {
 		return nil
 	}
 
+	gui.State.ContextManager.Lock()
+
 	if len(gui.State.ContextManager.ContextStack) == 0 {
 		gui.State.ContextManager.ContextStack = []types.Context{c}
 	} else {
@@ -57,6 +56,8 @@ func (gui *Gui) replaceContext(c types.Context) error {
 		gui.State.ContextManager.ContextStack = append(gui.State.ContextManager.ContextStack[0:len(gui.State.ContextManager.ContextStack)-1], c)
 	}
 
+	defer gui.State.ContextManager.Unlock()
+
 	return gui.activateContext(c)
 }
 
diff --git a/pkg/gui/context/list_context_trait.go b/pkg/gui/context/list_context_trait.go
index 6deb5dfc1..e508c8029 100644
--- a/pkg/gui/context/list_context_trait.go
+++ b/pkg/gui/context/list_context_trait.go
@@ -43,7 +43,7 @@ func (self *ListContextTrait) HandleFocus(opts ...types.OnFocusOpts) error {
 func (self *ListContextTrait) HandleFocusLost() error {
 	self.viewTrait.SetOriginX(0)
 
-	return self.Context.HandleFocus()
+	return self.Context.HandleFocusLost()
 }
 
 // OnFocus assumes that the content of the context has already been rendered to the view. OnRender is the function which actually renders the content to the view
diff --git a/pkg/gui/diffing.go b/pkg/gui/diffing.go
index a23772ab3..daa659f7f 100644
--- a/pkg/gui/diffing.go
+++ b/pkg/gui/diffing.go
@@ -33,15 +33,19 @@ func (gui *Gui) renderDiff() error {
 // which becomes an option when you bring up the diff menu, but when you're just
 // flicking through branches it will be using the local branch name.
 func (gui *Gui) currentDiffTerminals() []string {
-	switch gui.currentContext().GetKey() {
-	case "":
+	c := gui.currentSideContext()
+
+	if c.GetKey() == "" {
 		return nil
-	case context.FILES_CONTEXT_KEY, context.SUBMODULES_CONTEXT_KEY:
+	}
+
+	switch v := c.(type) {
+	case *context.WorkingTreeContext, *context.SubmodulesContext:
 		// TODO: should we just return nil here?
 		return []string{""}
-	case context.COMMIT_FILES_CONTEXT_KEY:
-		return []string{gui.State.Contexts.CommitFiles.GetRefName()}
-	case context.LOCAL_BRANCHES_CONTEXT_KEY:
+	case *context.CommitFilesContext:
+		return []string{v.GetRefName()}
+	case *context.BranchesContext:
 		// for our local branches we want to include both the branch and its upstream
 		branch := gui.State.Contexts.Branches.GetSelected()
 		if branch != nil {
@@ -52,13 +56,13 @@ func (gui *Gui) currentDiffTerminals() []string {
 			return names
 		}
 		return nil
-	default:
-		itemId := gui.getSideContextSelectedItemId()
-		if itemId == "" {
-			return nil
-		}
+	case types.IListContext:
+		itemId := v.GetSelectedItemId()
+
 		return []string{itemId}
 	}
+
+	return nil
 }
 
 func (gui *Gui) currentDiffTerminal() string {
diff --git a/pkg/gui/presentation/files_test.go b/pkg/gui/presentation/files_test.go
index 441ec7b69..fcca2c27d 100644
--- a/pkg/gui/presentation/files_test.go
+++ b/pkg/gui/presentation/files_test.go
@@ -70,6 +70,7 @@ M  file1
 		s := s
 		t.Run(s.name, func(t *testing.T) {
 			viewModel := filetree.NewFileTree(func() []*models.File { return s.files }, utils.NewDummyLog(), true)
+			viewModel.SetTree()
 			for _, path := range s.collapsedPaths {
 				viewModel.ToggleCollapsed(path)
 			}
@@ -128,6 +129,7 @@ M file1
 		s := s
 		t.Run(s.name, func(t *testing.T) {
 			viewModel := filetree.NewCommitFileTreeViewModel(func() []*models.CommitFile { return s.files }, utils.NewDummyLog(), true)
+			viewModel.SetTree()
 			for _, path := range s.collapsedPaths {
 				viewModel.ToggleCollapsed(path)
 			}