fix(init): Make sure to switch contexts on init

cmd/init.go

@@ -46,19 +46,35 @@ var initCmd = &cobra.Command{
 			rtOpts = []*runtime.Runtime{overridesVal.(*runtime.Runtime)}
 		}
 
+		contextName := "local"
+		changingContext := len(args) > 0
+		if changingContext {
+			contextName = args[0]
+
+			tempRt, err := runtime.NewRuntime(rtOpts...)
+			if err != nil {
+				return fmt.Errorf("failed to initialize context: %w", err)
+			}
+
+			if _, err := tempRt.Shell.WriteResetToken(); err != nil {
+				return fmt.Errorf("failed to write reset token: %w", err)
+			}
+
+			if err := tempRt.ConfigHandler.SetContext(contextName); err != nil {
+				return fmt.Errorf("failed to set context: %w", err)
+			}
+		}
+
 		rt, err := runtime.NewRuntime(rtOpts...)
 		if err != nil {
-			return fmt.Errorf("failed to initialize context: %w", err)
+			return fmt.Errorf("failed to initialize runtime: %w", err)
 		}
 
 		if err := rt.Shell.AddCurrentDirToTrustedFile(); err != nil {
 			return fmt.Errorf("failed to add current directory to trusted file: %w", err)
 		}
 
-		contextName := "local"
-		if len(args) > 0 {
-			contextName = args[0]
-		} else {
+		if !changingContext {
 			currentContext := rt.ConfigHandler.GetContext()
 			if currentContext != "" && currentContext != "local" {
 				contextName = currentContext
@@ -135,8 +151,10 @@ var initCmd = &cobra.Command{
 			return err
 		}
 
-		if err := rt.HandleSessionReset(); err != nil {
-			return fmt.Errorf("failed to handle session reset: %w", err)
+		if !changingContext {
+			if err := rt.HandleSessionReset(); err != nil {
+				return fmt.Errorf("failed to handle session reset: %w", err)
+			}
 		}
 
 		var blueprintURL []string
@@ -152,16 +170,6 @@ var initCmd = &cobra.Command{
 			return fmt.Errorf("failed to save configuration: %w", err)
 		}
 
-		if len(args) > 0 {
-			if _, err := rt.Shell.WriteResetToken(); err != nil {
-				return fmt.Errorf("failed to write reset token: %w", err)
-			}
-
-			if err := rt.ConfigHandler.SetContext(contextName); err != nil {
-				return fmt.Errorf("failed to set context: %w", err)
-			}
-		}
-
 		fmt.Fprintln(os.Stderr, "Initialization successful")
 
 		return nil

cmd/init_test.go

@@ -1081,7 +1081,7 @@ func TestInitCmd(t *testing.T) {
 		}
 	})
 
-	t.Run("CallsHandleSessionReset", func(t *testing.T) {
+	t.Run("CallsHandleSessionResetWhenNoContextProvided", func(t *testing.T) {
 		// Given a temporary directory with mocked dependencies
 		mocks := setupInitTest(t)
 
@@ -1092,7 +1092,7 @@ func TestInitCmd(t *testing.T) {
 			return false, nil
 		}
 
-		// When executing the init command
+		// When executing the init command without a context name
 		cmd := createTestInitCmd()
 		ctx := context.WithValue(context.Background(), runtimeOverridesKey, mocks.Runtime)
 		cmd.SetArgs([]string{})
@@ -1143,16 +1143,22 @@ func TestInitCmd(t *testing.T) {
 		// Given a temporary directory with mocked dependencies
 		mocks := setupInitTest(t)
 
-		// And tracking whether WriteResetToken and SetContext are called
+		// And tracking whether WriteResetToken, SetContext, and HandleSessionReset are called
 		var writeResetTokenCalled bool
 		var setContextCalled bool
 		var setContextValue string
+		var checkResetFlagsCalled bool
 
 		mocks.Shell.Shell.WriteResetTokenFunc = func() (string, error) {
 			writeResetTokenCalled = true
 			return "test-token", nil
 		}
 
+		mocks.Shell.Shell.CheckResetFlagsFunc = func() (bool, error) {
+			checkResetFlagsCalled = true
+			return false, nil
+		}
+
 		mockConfigHandler := mocks.ConfigHandler.(*config.MockConfigHandler)
 		mockConfigHandler.SetContextFunc = func(context string) error {
 			setContextCalled = true
@@ -1185,21 +1191,32 @@ func TestInitCmd(t *testing.T) {
 		if setContextValue != "test-context" {
 			t.Errorf("Expected SetContext to be called with 'test-context', got: %s", setContextValue)
 		}
+
+		// And HandleSessionReset should NOT have been called (skipped when changing contexts)
+		if checkResetFlagsCalled {
+			t.Error("Expected HandleSessionReset to not be called when context name is provided, but it was called")
+		}
 	})
 
 	t.Run("DoesNotSwitchContextWhenNoContextNameProvided", func(t *testing.T) {
 		// Given a temporary directory with mocked dependencies
 		mocks := setupInitTest(t)
 
-		// And tracking whether WriteResetToken and SetContext are called
+		// And tracking whether WriteResetToken, SetContext, and HandleSessionReset are called
 		var writeResetTokenCalled bool
 		var setContextCalled bool
+		var checkResetFlagsCalled bool
 
 		mocks.Shell.Shell.WriteResetTokenFunc = func() (string, error) {
 			writeResetTokenCalled = true
 			return "test-token", nil
 		}
 
+		mocks.Shell.Shell.CheckResetFlagsFunc = func() (bool, error) {
+			checkResetFlagsCalled = true
+			return false, nil
+		}
+
 		mockConfigHandler := mocks.ConfigHandler.(*config.MockConfigHandler)
 		mockConfigHandler.SetContextFunc = func(context string) error {
 			setContextCalled = true
@@ -1227,6 +1244,11 @@ func TestInitCmd(t *testing.T) {
 		if setContextCalled {
 			t.Error("Expected SetContext to not be called when no context name is provided, but it was called")
 		}
+
+		// And HandleSessionReset SHOULD have been called (not skipped when not changing contexts)
+		if !checkResetFlagsCalled {
+			t.Error("Expected HandleSessionReset to be called when no context name is provided, but it was not")
+		}
 	})
 
 	t.Run("HandlesSetContextError", func(t *testing.T) {