store: Have UpdateMachine.load check account still exists after async gaps

lib/model/store.dart

@@ -147,13 +147,17 @@ abstract class GlobalStore extends ChangeNotifier {
     // It's up to us.  Start loading.
     future = loadPerAccount(accountId);
     _perAccountStoresLoading[accountId] = future;
-    store = await future;
-    _setPerAccount(accountId, store);
-    unawaited(_perAccountStoresLoading.remove(accountId));
-    return store;
+    try {
+      store = await future;
+      _setPerAccount(accountId, store);
+      return store;
+    } finally {
+      unawaited(_perAccountStoresLoading.remove(accountId));
+    }
   }
 
   Future<void> _reloadPerAccount(int accountId) async {
+    assert(_accounts.containsKey(accountId));
     assert(_perAccountStores.containsKey(accountId));
     assert(!_perAccountStoresLoading.containsKey(accountId));
     final store = await loadPerAccount(accountId);
@@ -169,6 +173,11 @@ abstract class GlobalStore extends ChangeNotifier {
 
   /// Load per-account data for the given account, unconditionally.
   ///
+  /// The account for `accountId` must exist.
+  ///
+  /// Throws [AccountNotFoundException] if after any async gap
+  /// the account has been removed.
+  ///
   /// This method should be called only by the implementation of [perAccount].
   /// Other callers interested in per-account data should use [perAccount]
   /// and/or [perAccountSync].
@@ -183,36 +192,30 @@ abstract class GlobalStore extends ChangeNotifier {
           // The API key is invalid and the store can never be loaded
           // unless the user retries manually.
           final account = getAccount(accountId);
-          if (account == null) {
-            // The account was logged out during `await doLoadPerAccount`.
-            // Here, that seems possible only by the user's own action;
-            // the logout can't have been done programmatically.
-            // Even if it were, it would have come with its own UI feedback.
-            // Anyway, skip showing feedback, to not be confusing or repetitive.
-            throw AccountNotFoundException();
-          }
+          assert(account != null); // doLoadPerAccount would have thrown AccountNotFoundException
           final zulipLocalizations = GlobalLocalizations.zulipLocalizations;
           reportErrorToUserModally(
             zulipLocalizations.errorCouldNotConnectTitle,
             message: zulipLocalizations.errorInvalidApiKeyMessage(
-              account.realmUrl.toString()));
+              account!.realmUrl.toString()));
           await logOutAccount(this, accountId);
           throw AccountNotFoundException();
         default:
           rethrow;
       }
     }
-    if (!_accounts.containsKey(accountId)) {
-      // TODO(#1354): handle this earlier
-      // [removeAccount] was called during [doLoadPerAccount].
-      store.dispose();
-      throw AccountNotFoundException();
-    }
+    // doLoadPerAccount would have thrown AccountNotFoundException
+    assert(_accounts.containsKey(accountId));
     return store;
   }
 
   /// Load per-account data for the given account, unconditionally.
   ///
+  /// The account for `accountId` must exist.
+  ///
+  /// Throws [AccountNotFoundException] if after any async gap
+  /// the account has been removed.
+  ///
   /// This method should be called only by [loadPerAccount].
   Future<PerAccountStore> doLoadPerAccount(int accountId);
 
@@ -263,6 +266,8 @@ abstract class GlobalStore extends ChangeNotifier {
   Future<void> doUpdateAccount(int accountId, AccountsCompanion data);
 
   /// Remove an account from the store.
+  ///
+  /// The account for `accountId` must exist.
   Future<void> removeAccount(int accountId) async {
     assert(_accounts.containsKey(accountId));
     await doRemoveAccount(accountId);
@@ -941,15 +946,31 @@ class UpdateMachine {
     store.updateMachine = this;
   }
 
-  /// Load the user's data from the server, and start an event queue going.
+  /// Load data for the given account from the server,
+  /// and start an event queue going.
+  ///
+  /// The account for `accountId` must exist.
+  ///
+  /// Throws [AccountNotFoundException] if after any async gap
+  /// the account has been removed.
   ///
   /// In the future this might load an old snapshot from local storage first.
   static Future<UpdateMachine> load(GlobalStore globalStore, int accountId) async {
     Account account = globalStore.getAccount(accountId)!;
     final connection = globalStore.apiConnectionFromAccount(account);
 
+    void stopAndThrowIfNoAccount() {
+      final account = globalStore.getAccount(accountId);
+      if (account == null) {
+        assert(debugLog('Account logged out during UpdateMachine.load'));
+        connection.close();
+        throw AccountNotFoundException();
+      }
+    }
+
     final stopwatch = Stopwatch()..start();
-    final initialSnapshot = await _registerQueueWithRetry(connection);
+    final initialSnapshot = await _registerQueueWithRetry(connection,
+      stopAndThrowIfNoAccount: stopAndThrowIfNoAccount);
     final t = (stopwatch..stop()).elapsed;
     assert(debugLog("initial fetch time: ${t.inMilliseconds}ms"));
 
@@ -991,13 +1012,20 @@ class UpdateMachine {
 
   bool _disposed = false;
 
+  /// Make the register-queue request, with retries.
+  ///
+  /// After each async gap, calls [stopAndThrowIfNoAccount].
   static Future<InitialSnapshot> _registerQueueWithRetry(
-      ApiConnection connection) async {
+    ApiConnection connection, {
+    required void Function() stopAndThrowIfNoAccount,
+  }) async {
     BackoffMachine? backoffMachine;
     while (true) {
+      InitialSnapshot? result;
       try {
-        return await registerQueue(connection);
+        result = await registerQueue(connection);
       } catch (e, s) {
+        stopAndThrowIfNoAccount();
         // TODO(#890): tell user if initial-fetch errors persist, or look non-transient
         switch (e) {
           case HttpException(httpStatus: 401):
@@ -1012,8 +1040,13 @@ class UpdateMachine {
         }
         assert(debugLog('Backing off, then will retry…'));
         await (backoffMachine ??= BackoffMachine()).wait();
+        stopAndThrowIfNoAccount();
         assert(debugLog('… Backoff wait complete, retrying initial fetch.'));
       }
+      if (result != null) {
+        stopAndThrowIfNoAccount();
+        return result;
+      }
     }
   }
 

test/api/fake_api.dart

@@ -1,6 +1,7 @@
 import 'dart:collection';
 import 'dart:convert';
 
+import 'package:checks/checks.dart';
 import 'package:flutter/foundation.dart';
 import 'package:http/http.dart' as http;
 import 'package:zulip/api/core.dart';
@@ -278,3 +279,7 @@ class FakeApiConnection extends ApiConnection {
     );
   }
 }
+
+extension FakeApiConnectionChecks on Subject<FakeApiConnection> {
+  Subject<bool> get isOpen => has((x) => x.isOpen, 'isOpen');
+}

test/example_data.dart

@@ -981,9 +981,14 @@ PerAccountStore store({
 }
 const _store = store;
 
-UpdateMachine updateMachine({Account? account, InitialSnapshot? initialSnapshot}) {
+UpdateMachine updateMachine({
+  GlobalStore? globalStore,
+  Account? account,
+  InitialSnapshot? initialSnapshot,
+}) {
   initialSnapshot ??= _initialSnapshot();
-  final store = _store(account: account, initialSnapshot: initialSnapshot);
+  final store = _store(globalStore: globalStore,
+    account: account, initialSnapshot: initialSnapshot);
   return UpdateMachine.fromInitialSnapshot(
     store: store, initialSnapshot: initialSnapshot);
 }