Snap for 8712302 from 5a0dcd8114afa3b0f93bf77a6c989f539555186c to tm-frc-art-releaset_frc_art_330443060android13-frc-art-release

Change-Id: I8de253b6b6d24906b00b6a63e35f6ee221192c28

12 files changed, 496 insertions, 125 deletions

diff --git a/icing/file/file-backed-proto.h b/icing/file/file-backed-proto.h
index 15a1953..d7d9bad 100644
--- a/icing/file/file-backed-proto.h
+++ b/icing/file/file-backed-proto.h
@@ -63,6 +63,17 @@ class FileBackedProto {
   // file_path : Must be a path within in a directory that already exists.
   FileBackedProto(const Filesystem& filesystem, std::string_view file_path);
 
+  // Reset the internal file_path for the file backed proto.
+  // Example use:
+  //   auto file_backed_proto1 = *FileBackedProto<Proto>::Create(...);
+  //   auto file_backed_proto2 = *FileBackedProto<Proto>::Create(...);
+  //   filesystem.SwapFiles(file1, file2);
+  //   file_backed_proto1.SetSwappedFilepath(file2);
+  //   file_backed_proto2.SetSwappedFilepath(file1);
+  void SetSwappedFilepath(std::string_view swapped_to_file_path) {
+    file_path_ = swapped_to_file_path;
+  }
+
   // Returns a reference to the proto read from the file. It
   // internally caches the read proto so that future calls are fast.
   //
@@ -99,7 +110,7 @@ class FileBackedProto {
   mutable absl_ports::shared_mutex mutex_;
 
   const Filesystem* const filesystem_;
-  const std::string file_path_;
+  std::string file_path_;
 
   mutable std::unique_ptr<ProtoT> cached_proto_ ICING_GUARDED_BY(mutex_);
 };
diff --git a/icing/file/file-backed-vector.h b/icing/file/file-backed-vector.h
index 00bdc7e..7e42e32 100644
--- a/icing/file/file-backed-vector.h
+++ b/icing/file/file-backed-vector.h
@@ -586,8 +586,11 @@ libtextclassifier3::Status FileBackedVector<T>::GrowIfNecessary(
   }
 
   int64_t current_file_size = filesystem_->GetFileSize(file_path_.c_str());
-  int64_t least_file_size_needed = sizeof(Header) + num_elements * sizeof(T);
+  if (current_file_size == Filesystem::kBadFileSize) {
+    return absl_ports::InternalError("Unable to retrieve file size.");
+  }
 
+  int64_t least_file_size_needed = sizeof(Header) + num_elements * sizeof(T);
   if (least_file_size_needed <= current_file_size) {
     // Our underlying file can hold the target num_elements cause we've grown
     // before
diff --git a/icing/file/file-backed-vector_test.cc b/icing/file/file-backed-vector_test.cc
index 54f9ef5..ed94fa5 100644
--- a/icing/file/file-backed-vector_test.cc
+++ b/icing/file/file-backed-vector_test.cc
@@ -677,6 +677,27 @@ TEST_F(FileBackedVectorTest, RemapFailureStillValidInstance) {
   EXPECT_THAT(vector->Get(kResizingIndex / 2), IsOkAndHolds(Pointee(Eq(9))));
 }
 
+TEST_F(FileBackedVectorTest, BadFileSizeDuringGrowReturnsError) {
+  auto mock_filesystem = std::make_unique<MockFilesystem>();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<FileBackedVector<int>> vector,
+      FileBackedVector<int>::Create(
+          *mock_filesystem, file_path_,
+          MemoryMappedFile::Strategy::READ_WRITE_AUTO_SYNC));
+
+  // At first, the vector is empty and has no mapping established. The first Set
+  // call will cause a Grow.
+  // During Grow, we will attempt to check the underlying file size to see if
+  // growing is actually necessary. Return an error on the call to GetFileSize.
+  ON_CALL(*mock_filesystem, GetFileSize(A<const char*>()))
+      .WillByDefault(Return(Filesystem::kBadFileSize));
+
+  // We should fail gracefully and return an INTERNAL error to indicate that
+  // there was an issue retrieving the file size.
+  EXPECT_THAT(vector->Set(0, 7),
+              StatusIs(libtextclassifier3::StatusCode::INTERNAL));
+}
+
 }  // namespace
 
 }  // namespace lib
diff --git a/icing/icing-search-engine_test.cc b/icing/icing-search-engine_test.cc
index 5244f4c..13e77b8 100644
--- a/icing/icing-search-engine_test.cc
+++ b/icing/icing-search-engine_test.cc
@@ -765,8 +765,7 @@ TEST_F(IcingSearchEngineTest, FailToWriteSchema) {
 
   auto mock_filesystem = std::make_unique<MockFilesystem>();
   // This fails FileBackedProto::Write()
-  ON_CALL(*mock_filesystem,
-          OpenForWrite(Eq(icing_options.base_dir() + "/schema_dir/schema.pb")))
+  ON_CALL(*mock_filesystem, OpenForWrite(HasSubstr("schema.pb")))
       .WillByDefault(Return(-1));
 
   TestIcingSearchEngine icing(icing_options, std::move(mock_filesystem),
@@ -2993,13 +2992,17 @@ TEST_F(IcingSearchEngineTest, OptimizationFailureUninitializesIcing) {
   };
   ON_CALL(*mock_filesystem, CreateDirectoryRecursively)
       .WillByDefault(create_dir_lambda);
+
   auto swap_lambda = [&just_swapped_files](const char* first_dir,
                                            const char* second_dir) {
     just_swapped_files = true;
     return false;
   };
-  ON_CALL(*mock_filesystem, SwapFiles).WillByDefault(swap_lambda);
-  TestIcingSearchEngine icing(GetDefaultIcingOptions(),
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  ON_CALL(*mock_filesystem, SwapFiles(HasSubstr("document_dir_optimize_tmp"),
+                                      HasSubstr("document_dir")))
+      .WillByDefault(swap_lambda);
+  TestIcingSearchEngine icing(options, std::move(mock_filesystem),
                               std::move(mock_filesystem),
                               std::make_unique<IcingFilesystem>(),
                               std::make_unique<FakeClock>(), GetTestJniCache());
@@ -3752,7 +3755,8 @@ TEST_F(IcingSearchEngineTest, IcingShouldWorkFineIfOptimizationIsAborted) {
   // fails. This will fail IcingSearchEngine::OptimizeDocumentStore() and makes
   // it return ABORTED_ERROR.
   auto mock_filesystem = std::make_unique<MockFilesystem>();
-  ON_CALL(*mock_filesystem, DeleteDirectoryRecursively)
+  ON_CALL(*mock_filesystem,
+          DeleteDirectoryRecursively(HasSubstr("_optimize_tmp")))
       .WillByDefault(Return(false));
 
   TestIcingSearchEngine icing(GetDefaultIcingOptions(),
@@ -3799,7 +3803,8 @@ TEST_F(IcingSearchEngineTest,
   // Creates a mock filesystem in which SwapFiles() always fails and deletes the
   // directories. This will fail IcingSearchEngine::OptimizeDocumentStore().
   auto mock_filesystem = std::make_unique<MockFilesystem>();
-  ON_CALL(*mock_filesystem, SwapFiles)
+  ON_CALL(*mock_filesystem, SwapFiles(HasSubstr("document_dir_optimize_tmp"),
+                                      HasSubstr("document_dir")))
       .WillByDefault([this](const char* one, const char* two) {
         filesystem()->DeleteDirectoryRecursively(one);
         filesystem()->DeleteDirectoryRecursively(two);
@@ -3870,7 +3875,8 @@ TEST_F(IcingSearchEngineTest, OptimizationShouldRecoverIfDataFilesAreMissing) {
   // Creates a mock filesystem in which SwapFiles() always fails and empties the
   // directories. This will fail IcingSearchEngine::OptimizeDocumentStore().
   auto mock_filesystem = std::make_unique<MockFilesystem>();
-  ON_CALL(*mock_filesystem, SwapFiles)
+  ON_CALL(*mock_filesystem, SwapFiles(HasSubstr("document_dir_optimize_tmp"),
+                                      HasSubstr("document_dir")))
       .WillByDefault([this](const char* one, const char* two) {
         filesystem()->DeleteDirectoryRecursively(one);
         filesystem()->CreateDirectoryRecursively(one);
diff --git a/icing/index/index-processor_test.cc b/icing/index/index-processor_test.cc
index bd310de..7746688 100644
--- a/icing/index/index-processor_test.cc
+++ b/icing/index/index-processor_test.cc
@@ -153,9 +153,12 @@ class IndexProcessorTest : public Test {
         normalizer_factory::Create(
             /*max_term_byte_size=*/std::numeric_limits<int32_t>::max()));
 
+    std::string schema_store_dir = GetTestTempDir() + "/schema_store";
+    ASSERT_TRUE(
+        filesystem_.CreateDirectoryRecursively(schema_store_dir.c_str()));
     ICING_ASSERT_OK_AND_ASSIGN(
         schema_store_,
-        SchemaStore::Create(&filesystem_, GetTestTempDir(), &fake_clock_));
+        SchemaStore::Create(&filesystem_, schema_store_dir, &fake_clock_));
     SchemaProto schema =
         SchemaBuilder()
             .AddType(
diff --git a/icing/query/query-processor_test.cc b/icing/query/query-processor_test.cc
index 950f739..eaa0efc 100644
--- a/icing/query/query-processor_test.cc
+++ b/icing/query/query-processor_test.cc
@@ -77,12 +77,14 @@ class QueryProcessorTest : public Test {
   QueryProcessorTest()
       : test_dir_(GetTestTempDir() + "/icing"),
         store_dir_(test_dir_ + "/store"),
+        schema_store_dir_(test_dir_ + "/schema_store"),
         index_dir_(test_dir_ + "/index") {}
 
   void SetUp() override {
     filesystem_.DeleteDirectoryRecursively(test_dir_.c_str());
     filesystem_.CreateDirectoryRecursively(index_dir_.c_str());
     filesystem_.CreateDirectoryRecursively(store_dir_.c_str());
+    filesystem_.CreateDirectoryRecursively(schema_store_dir_.c_str());
 
     if (!IsCfStringTokenization() && !IsReverseJniTokenization()) {
       // If we've specified using the reverse-JNI method for segmentation (i.e.
@@ -129,6 +131,7 @@ class QueryProcessorTest : public Test {
   Filesystem filesystem_;
   const std::string test_dir_;
   const std::string store_dir_;
+  const std::string schema_store_dir_;
   std::unique_ptr<Index> index_;
   std::unique_ptr<LanguageSegmenter> language_segmenter_;
   std::unique_ptr<Normalizer> normalizer_;
@@ -176,7 +179,7 @@ TEST_F(QueryProcessorTest, EmptyGroupMatchAllDocuments) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -227,7 +230,7 @@ TEST_F(QueryProcessorTest, EmptyQueryMatchAllDocuments) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -278,7 +281,7 @@ TEST_F(QueryProcessorTest, QueryTermNormalized) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -354,7 +357,7 @@ TEST_F(QueryProcessorTest, OneTermPrefixMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -422,7 +425,7 @@ TEST_F(QueryProcessorTest, OneTermExactMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -490,7 +493,7 @@ TEST_F(QueryProcessorTest, AndSameTermExactMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -560,7 +563,7 @@ TEST_F(QueryProcessorTest, AndTwoTermExactMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -635,7 +638,7 @@ TEST_F(QueryProcessorTest, AndSameTermPrefixMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -705,7 +708,7 @@ TEST_F(QueryProcessorTest, AndTwoTermPrefixMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -781,7 +784,7 @@ TEST_F(QueryProcessorTest, AndTwoTermPrefixAndExactMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -857,7 +860,7 @@ TEST_F(QueryProcessorTest, OrTwoTermExactMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -946,7 +949,7 @@ TEST_F(QueryProcessorTest, OrTwoTermPrefixMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1034,7 +1037,7 @@ TEST_F(QueryProcessorTest, OrTwoTermPrefixAndExactMatch) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1121,7 +1124,7 @@ TEST_F(QueryProcessorTest, CombinedAndOrTerms) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1307,7 +1310,7 @@ TEST_F(QueryProcessorTest, OneGroup) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1383,7 +1386,7 @@ TEST_F(QueryProcessorTest, TwoGroups) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1461,7 +1464,7 @@ TEST_F(QueryProcessorTest, ManyLevelNestedGrouping) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1537,7 +1540,7 @@ TEST_F(QueryProcessorTest, OneLevelNestedGrouping) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1614,7 +1617,7 @@ TEST_F(QueryProcessorTest, ExcludeTerm) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1679,7 +1682,7 @@ TEST_F(QueryProcessorTest, ExcludeNonexistentTerm) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1742,7 +1745,7 @@ TEST_F(QueryProcessorTest, ExcludeAnd) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1832,7 +1835,7 @@ TEST_F(QueryProcessorTest, ExcludeOr) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1928,7 +1931,7 @@ TEST_F(QueryProcessorTest, DeletedFilter) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2002,7 +2005,7 @@ TEST_F(QueryProcessorTest, NamespaceFilter) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2078,7 +2081,7 @@ TEST_F(QueryProcessorTest, SchemaTypeFilter) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2155,7 +2158,7 @@ TEST_F(QueryProcessorTest, SectionFilterForOneDocument) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2237,7 +2240,7 @@ TEST_F(QueryProcessorTest, SectionFilterAcrossSchemaTypes) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2320,7 +2323,7 @@ TEST_F(QueryProcessorTest, SectionFilterWithinSchemaType) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2404,7 +2407,7 @@ TEST_F(QueryProcessorTest, SectionFilterRespectsDifferentSectionIds) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2477,7 +2480,7 @@ TEST_F(QueryProcessorTest, NonexistentSectionFilterReturnsEmptyResults) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2544,7 +2547,7 @@ TEST_F(QueryProcessorTest, UnindexedSectionFilterReturnsEmptyResults) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2614,7 +2617,7 @@ TEST_F(QueryProcessorTest, SectionFilterTermAndUnrestrictedTerm) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -2689,7 +2692,7 @@ TEST_F(QueryProcessorTest, DocumentBeforeTtlNotFilteredOut) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   // Arbitrary value, just has to be less than the document's creation
@@ -2748,7 +2751,7 @@ TEST_F(QueryProcessorTest, DocumentPastTtlFilteredOut) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       schema_store_,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
   // Arbitrary value, just has to be greater than the document's creation
diff --git a/icing/schema/schema-store.cc b/icing/schema/schema-store.cc
index acc5030..fc50ea6 100644
--- a/icing/schema/schema-store.cc
+++ b/icing/schema/schema-store.cc
@@ -108,27 +108,60 @@ libtextclassifier3::StatusOr<std::unique_ptr<SchemaStore>> SchemaStore::Create(
   ICING_RETURN_ERROR_IF_NULL(filesystem);
   ICING_RETURN_ERROR_IF_NULL(clock);
 
+  if (!filesystem->DirectoryExists(base_dir.c_str())) {
+    return absl_ports::FailedPreconditionError(
+        "Schema store base directory does not exist!");
+  }
   std::unique_ptr<SchemaStore> schema_store = std::unique_ptr<SchemaStore>(
       new SchemaStore(filesystem, base_dir, clock));
   ICING_RETURN_IF_ERROR(schema_store->Initialize(initialize_stats));
   return schema_store;
 }
 
+libtextclassifier3::StatusOr<std::unique_ptr<SchemaStore>> SchemaStore::Create(
+    const Filesystem* filesystem, const std::string& base_dir,
+    const Clock* clock, SchemaProto schema) {
+  ICING_RETURN_ERROR_IF_NULL(filesystem);
+  ICING_RETURN_ERROR_IF_NULL(clock);
+
+  if (!filesystem->DirectoryExists(base_dir.c_str())) {
+    return absl_ports::FailedPreconditionError(
+        "Schema store base directory does not exist!");
+  }
+  std::unique_ptr<SchemaStore> schema_store = std::unique_ptr<SchemaStore>(
+      new SchemaStore(filesystem, base_dir, clock));
+  ICING_RETURN_IF_ERROR(schema_store->Initialize(std::move(schema)));
+  return schema_store;
+}
+
 SchemaStore::SchemaStore(const Filesystem* filesystem, std::string base_dir,
                          const Clock* clock)
-    : filesystem_(*filesystem),
+    : filesystem_(filesystem),
       base_dir_(std::move(base_dir)),
-      clock_(*clock),
-      schema_file_(*filesystem, MakeSchemaFilename(base_dir_)) {}
+      clock_(clock),
+      schema_file_(std::make_unique<FileBackedProto<SchemaProto>>(
+          *filesystem, MakeSchemaFilename(base_dir_))) {}
 
 SchemaStore::~SchemaStore() {
-  if (has_schema_successfully_set_) {
+  if (has_schema_successfully_set_ && schema_file_ != nullptr &&
+      schema_type_mapper_ != nullptr && section_manager_ != nullptr) {
     if (!PersistToDisk().ok()) {
       ICING_LOG(ERROR) << "Error persisting to disk in SchemaStore destructor";
     }
   }
 }
 
+libtextclassifier3::Status SchemaStore::Initialize(SchemaProto new_schema) {
+  if (!absl_ports::IsNotFound(GetSchema().status())) {
+    return absl_ports::FailedPreconditionError(
+        "Incorrectly tried to initialize schema store with a new schema, when "
+        "one is already set!");
+  }
+  ICING_RETURN_IF_ERROR(schema_file_->Write(
+      std::make_unique<SchemaProto>(std::move(new_schema))));
+  return InitializeInternal(/*initialize_stats=*/nullptr);
+}
+
 libtextclassifier3::Status SchemaStore::Initialize(
     InitializeStatsProto* initialize_stats) {
   auto schema_proto_or = GetSchema();
@@ -139,13 +172,16 @@ libtextclassifier3::Status SchemaStore::Initialize(
     // Real error when trying to read the existing schema
     return schema_proto_or.status();
   }
-  has_schema_successfully_set_ = true;
+  return InitializeInternal(initialize_stats);
+}
 
+libtextclassifier3::Status SchemaStore::InitializeInternal(
+    InitializeStatsProto* initialize_stats) {
   if (!InitializeDerivedFiles().ok()) {
     ICING_VLOG(3)
         << "Couldn't find derived files or failed to initialize them, "
            "regenerating derived files for SchemaStore.";
-    std::unique_ptr<Timer> regenerate_timer = clock_.GetNewTimer();
+    std::unique_ptr<Timer> regenerate_timer = clock_->GetNewTimer();
     if (initialize_stats != nullptr) {
       initialize_stats->set_schema_store_recovery_cause(
           InitializeStatsProto::IO_ERROR);
@@ -161,6 +197,7 @@ libtextclassifier3::Status SchemaStore::Initialize(
     initialize_stats->set_num_schema_types(type_config_map_.size());
   }
 
+  has_schema_successfully_set_ = true;
   return libtextclassifier3::Status::OK;
 }
 
@@ -172,8 +209,8 @@ libtextclassifier3::Status SchemaStore::InitializeDerivedFiles() {
   }
 
   SchemaStore::Header header;
-  if (!filesystem_.Read(MakeHeaderFilename(base_dir_).c_str(), &header,
-                        sizeof(header))) {
+  if (!filesystem_->Read(MakeHeaderFilename(base_dir_).c_str(), &header,
+                         sizeof(header))) {
     return absl_ports::InternalError(
         absl_ports::StrCat("Couldn't read: ", MakeHeaderFilename(base_dir_)));
   }
@@ -185,7 +222,7 @@ libtextclassifier3::Status SchemaStore::InitializeDerivedFiles() {
 
   ICING_ASSIGN_OR_RETURN(
       schema_type_mapper_,
-      KeyMapper<SchemaTypeId>::Create(filesystem_,
+      KeyMapper<SchemaTypeId>::Create(*filesystem_,
                                       MakeSchemaTypeMapperFilename(base_dir_),
                                       kSchemaTypeMapperMaxSize));
 
@@ -236,12 +273,12 @@ libtextclassifier3::Status SchemaStore::RegenerateDerivedFiles() {
 }
 
 bool SchemaStore::HeaderExists() {
-  if (!filesystem_.FileExists(MakeHeaderFilename(base_dir_).c_str())) {
+  if (!filesystem_->FileExists(MakeHeaderFilename(base_dir_).c_str())) {
     return false;
   }
 
   int64_t file_size =
-      filesystem_.GetFileSize(MakeHeaderFilename(base_dir_).c_str());
+      filesystem_->GetFileSize(MakeHeaderFilename(base_dir_).c_str());
 
   // If it's been truncated to size 0 before, we consider it to be a new file
   return file_size != 0 && file_size != Filesystem::kBadFileSize;
@@ -254,11 +291,11 @@ libtextclassifier3::Status SchemaStore::UpdateHeader(const Crc32& checksum) {
   header.checksum = checksum.Get();
 
   ScopedFd scoped_fd(
-      filesystem_.OpenForWrite(MakeHeaderFilename(base_dir_).c_str()));
+      filesystem_->OpenForWrite(MakeHeaderFilename(base_dir_).c_str()));
   // This should overwrite the header.
   if (!scoped_fd.is_valid() ||
-      !filesystem_.Write(scoped_fd.get(), &header, sizeof(header)) ||
-      !filesystem_.DataSync(scoped_fd.get())) {
+      !filesystem_->Write(scoped_fd.get(), &header, sizeof(header)) ||
+      !filesystem_->DataSync(scoped_fd.get())) {
     return absl_ports::InternalError(absl_ports::StrCat(
         "Failed to write SchemaStore header: ", MakeHeaderFilename(base_dir_)));
   }
@@ -271,7 +308,7 @@ libtextclassifier3::Status SchemaStore::ResetSchemaTypeMapper() {
   // TODO(b/216487496): Implement a more robust version of TC_RETURN_IF_ERROR
   // that can support error logging.
   libtextclassifier3::Status status = KeyMapper<SchemaTypeId>::Delete(
-      filesystem_, MakeSchemaTypeMapperFilename(base_dir_));
+      *filesystem_, MakeSchemaTypeMapperFilename(base_dir_));
   if (!status.ok()) {
     ICING_LOG(ERROR) << status.error_message()
                      << "Failed to delete old schema_type mapper";
@@ -279,7 +316,7 @@ libtextclassifier3::Status SchemaStore::ResetSchemaTypeMapper() {
   }
   ICING_ASSIGN_OR_RETURN(
       schema_type_mapper_,
-      KeyMapper<SchemaTypeId>::Create(filesystem_,
+      KeyMapper<SchemaTypeId>::Create(*filesystem_,
                                       MakeSchemaTypeMapperFilename(base_dir_),
                                       kSchemaTypeMapperMaxSize));
 
@@ -287,17 +324,17 @@ libtextclassifier3::Status SchemaStore::ResetSchemaTypeMapper() {
 }
 
 libtextclassifier3::StatusOr<Crc32> SchemaStore::ComputeChecksum() const {
-  Crc32 total_checksum;
-  if (!has_schema_successfully_set_) {
-    // Nothing to checksum
-    return total_checksum;
+  auto schema_proto_or = GetSchema();
+  if (absl_ports::IsNotFound(schema_proto_or.status())) {
+    return Crc32();
   }
-  ICING_ASSIGN_OR_RETURN(const SchemaProto* schema_proto, GetSchema());
+  ICING_ASSIGN_OR_RETURN(const SchemaProto* schema_proto, schema_proto_or);
   Crc32 schema_checksum;
   schema_checksum.Append(schema_proto->SerializeAsString());
 
   Crc32 schema_type_mapper_checksum = schema_type_mapper_->ComputeChecksum();
 
+  Crc32 total_checksum;
   total_checksum.Append(std::to_string(schema_checksum.Get()));
   total_checksum.Append(std::to_string(schema_type_mapper_checksum.Get()));
 
@@ -306,7 +343,7 @@ libtextclassifier3::StatusOr<Crc32> SchemaStore::ComputeChecksum() const {
 
 libtextclassifier3::StatusOr<const SchemaProto*> SchemaStore::GetSchema()
     const {
-  return schema_file_.Read();
+  return schema_file_->Read();
 }
 
 // TODO(cassiewang): Consider removing this definition of SetSchema if it's not
@@ -393,17 +430,80 @@ SchemaStore::SetSchema(SchemaProto&& new_schema,
   result.success = result.success || ignore_errors_and_delete_documents;
 
   if (result.success) {
-    // Write the schema (and potentially overwrite a previous schema)
-    ICING_RETURN_IF_ERROR(
-        schema_file_.Write(std::make_unique<SchemaProto>(new_schema)));
+    ICING_RETURN_IF_ERROR(ApplySchemaChange(std::move(new_schema)));
     has_schema_successfully_set_ = true;
-
-    ICING_RETURN_IF_ERROR(RegenerateDerivedFiles());
   }
 
   return result;
 }
 
+libtextclassifier3::Status SchemaStore::ApplySchemaChange(
+    SchemaProto new_schema) {
+  // We need to ensure that we either 1) successfully set the schema and
+  // update all derived data structures or 2) fail and leave the schema store
+  // unchanged.
+  // So, first, we create an empty temporary directory to build a new schema
+  // store in.
+  std::string temp_schema_store_dir_path = base_dir_ + "_temp";
+  if (!filesystem_->DeleteDirectoryRecursively(
+          temp_schema_store_dir_path.c_str())) {
+    ICING_LOG(WARNING) << "Failed to recursively delete "
+                     << temp_schema_store_dir_path.c_str();
+    return absl_ports::InternalError(
+        "Unable to delete temp directory to prepare to build new schema "
+        "store.");
+  }
+
+  if (!filesystem_->CreateDirectoryRecursively(
+      temp_schema_store_dir_path.c_str())) {
+    return absl_ports::InternalError(
+        "Unable to create temp directory to build new schema store.");
+  }
+
+  // Then we create our new schema store with the new schema.
+  auto new_schema_store_or =
+      SchemaStore::Create(filesystem_, temp_schema_store_dir_path, clock_,
+                          std::move(new_schema));
+  if (!new_schema_store_or.ok()) {
+    // Attempt to clean up the temp directory.
+    if (!filesystem_->DeleteDirectoryRecursively(
+            temp_schema_store_dir_path.c_str())) {
+      // Nothing to do here. Just log an error.
+      ICING_LOG(WARNING) << "Failed to recursively delete "
+                       << temp_schema_store_dir_path.c_str();
+    }
+    return new_schema_store_or.status();
+  }
+  std::unique_ptr<SchemaStore> new_schema_store =
+      std::move(new_schema_store_or).ValueOrDie();
+
+  // Then we swap the new schema file + new derived files with the old files.
+  if (!filesystem_->SwapFiles(base_dir_.c_str(),
+                              temp_schema_store_dir_path.c_str())) {
+    // Attempt to clean up the temp directory.
+    if (!filesystem_->DeleteDirectoryRecursively(
+            temp_schema_store_dir_path.c_str())) {
+      // Nothing to do here. Just log an error.
+      ICING_LOG(WARNING) << "Failed to recursively delete "
+                       << temp_schema_store_dir_path.c_str();
+    }
+    return absl_ports::InternalError(
+        "Unable to apply new schema due to failed swap!");
+  }
+
+  std::string old_base_dir = std::move(base_dir_);
+  *this = std::move(*new_schema_store);
+
+  // After the std::move, the filepaths saved in this instance and in the
+  // schema_file_ instance will still be the one from temp_schema_store_dir
+  // even though they now point to files that are within old_base_dir.
+  // Manually set them to the correct paths.
+  base_dir_ = std::move(old_base_dir);
+  schema_file_->SetSwappedFilepath(MakeSchemaFilename(base_dir_));
+
+  return libtextclassifier3::Status::OK;
+}
+
 libtextclassifier3::StatusOr<const SchemaTypeConfigProto*>
 SchemaStore::GetSchemaTypeConfig(std::string_view schema_type) const {
   ICING_RETURN_IF_ERROR(CheckSchemaSet());
@@ -463,7 +563,7 @@ libtextclassifier3::Status SchemaStore::PersistToDisk() {
 
 SchemaStoreStorageInfoProto SchemaStore::GetStorageInfo() const {
   SchemaStoreStorageInfoProto storage_info;
-  int64_t directory_size = filesystem_.GetDiskUsage(base_dir_.c_str());
+  int64_t directory_size = filesystem_->GetDiskUsage(base_dir_.c_str());
   storage_info.set_schema_store_size(
       Filesystem::SanitizeFileSize(directory_size));
   ICING_ASSIGN_OR_RETURN(const SchemaProto* schema, GetSchema(), storage_info);
diff --git a/icing/schema/schema-store.h b/icing/schema/schema-store.h
index 2d3aca7..58e5477 100644
--- a/icing/schema/schema-store.h
+++ b/icing/schema/schema-store.h
@@ -130,8 +130,10 @@ class SchemaStore {
   static libtextclassifier3::StatusOr<std::unique_ptr<SchemaStore>> Create(
       const Filesystem* filesystem, const std::string& base_dir,
       const Clock* clock, InitializeStatsProto* initialize_stats = nullptr);
+  
+  SchemaStore(SchemaStore&&) = default;
+  SchemaStore& operator=(SchemaStore&&) = default;
 
-  // Not copyable
   SchemaStore(const SchemaStore&) = delete;
   SchemaStore& operator=(const SchemaStore&) = delete;
 
@@ -265,17 +267,51 @@ class SchemaStore {
   libtextclassifier3::StatusOr<SchemaDebugInfoProto> GetDebugInfo() const;
 
  private:
+  // Factory function to create a SchemaStore and set its schema. The created
+  // instance does not take ownership of any input components and all pointers
+  // must refer to valid objects that outlive the created SchemaStore instance.
+  // The base_dir must already exist. No schema must have set in base_dir prior
+  // to this.
+  //
+  // Returns:
+  //   A SchemaStore on success
+  //   FAILED_PRECONDITION on any null pointer input or if there has already
+  //       been a schema set for this path.
+  //   INTERNAL_ERROR on any IO errors
+  static libtextclassifier3::StatusOr<std::unique_ptr<SchemaStore>> Create(
+      const Filesystem* filesystem, const std::string& base_dir,
+      const Clock* clock, SchemaProto schema);
+
+
   // Use SchemaStore::Create instead.
   explicit SchemaStore(const Filesystem* filesystem, std::string base_dir,
                        const Clock* clock);
 
-  // Handles initializing the SchemaStore and regenerating any data if needed.
+  // Verifies that there is no error retrieving a previously set schema. Then
+  // initializes like normal.
   //
   // Returns:
   //   OK on success
   //   INTERNAL_ERROR on IO error
   libtextclassifier3::Status Initialize(InitializeStatsProto* initialize_stats);
 
+  // First, blindly writes new_schema to the schema_file. Then initializes like
+  // normal.
+  //
+  // Returns:
+  //   OK on success
+  //   INTERNAL_ERROR on IO error
+  //   FAILED_PRECONDITION if there is already a schema set for the schema_file.
+  libtextclassifier3::Status Initialize(SchemaProto new_schema);
+
+  // Handles initializing the SchemaStore and regenerating any data if needed.
+  //
+  // Returns:
+  //   OK on success
+  //   INTERNAL_ERROR on IO error
+  libtextclassifier3::Status InitializeInternal(
+      InitializeStatsProto* initialize_stats);
+
   // Creates sub-components and verifies the integrity of each sub-component.
   //
   // Returns:
@@ -310,15 +346,25 @@ class SchemaStore {
   // Returns any IO errors.
   libtextclassifier3::Status ResetSchemaTypeMapper();
 
+  // Creates a new schema store with new_schema and then swaps that new schema
+  // store with the existing one. This function guarantees that either: this
+  // instance will be fully updated to the new schema or no changes will take
+  // effect.
+  //
+  // Returns:
+  //   OK on success
+  //   INTERNAL on I/O error.
+  libtextclassifier3::Status ApplySchemaChange(SchemaProto new_schema);
+
   libtextclassifier3::Status CheckSchemaSet() const {
     return has_schema_successfully_set_
                ? libtextclassifier3::Status::OK
                : absl_ports::FailedPreconditionError("Schema not set yet.");
   }
 
-  const Filesystem& filesystem_;
-  const std::string base_dir_;
-  const Clock& clock_;
+  const Filesystem* filesystem_;
+  std::string base_dir_;
+  const Clock* clock_;
 
   // Used internally to indicate whether the class has been successfully
   // initialized with a valid schema. Will be false if Initialize failed or no
@@ -326,7 +372,7 @@ class SchemaStore {
   bool has_schema_successfully_set_ = false;
 
   // Cached schema
-  FileBackedProto<SchemaProto> schema_file_;
+  std::unique_ptr<FileBackedProto<SchemaProto>> schema_file_;
 
   // A hash map of (type config name -> type config), allows faster lookup of
   // type config in schema. The O(1) type config access makes schema-related and
diff --git a/icing/schema/schema-store_test.cc b/icing/schema/schema-store_test.cc
index 541918f..3fd41c4 100644
--- a/icing/schema/schema-store_test.cc
+++ b/icing/schema/schema-store_test.cc
@@ -21,7 +21,9 @@
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "icing/absl_ports/str_cat.h"
+#include "icing/document-builder.h"
 #include "icing/file/filesystem.h"
+#include "icing/file/mock-filesystem.h"
 #include "icing/portable/equals-proto.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
@@ -33,6 +35,7 @@
 #include "icing/testing/common-matchers.h"
 #include "icing/testing/fake-clock.h"
 #include "icing/testing/tmp-directory.h"
+#include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/util/crc32.h"
 
 namespace icing {
@@ -45,8 +48,11 @@ using ::testing::ElementsAre;
 using ::testing::Eq;
 using ::testing::Ge;
 using ::testing::Gt;
+using ::testing::HasSubstr;
 using ::testing::Not;
 using ::testing::Pointee;
+using ::testing::Return;
+using ::testing::SizeIs;
 
 constexpr PropertyConfigProto::Cardinality::Code CARDINALITY_OPTIONAL =
     PropertyConfigProto::Cardinality::OPTIONAL;
@@ -66,8 +72,10 @@ constexpr PropertyConfigProto::DataType::Code TYPE_DOUBLE =
 
 class SchemaStoreTest : public ::testing::Test {
  protected:
-  SchemaStoreTest() : test_dir_(GetTestTempDir() + "/icing") {
-    filesystem_.CreateDirectoryRecursively(test_dir_.c_str());
+  void SetUp() override {
+    temp_dir_ = GetTestTempDir() + "/icing";
+    schema_store_dir_ = temp_dir_ + "/schema_store";
+    filesystem_.CreateDirectoryRecursively(schema_store_dir_.c_str());
 
     schema_ =
         SchemaBuilder()
@@ -81,26 +89,112 @@ class SchemaStoreTest : public ::testing::Test {
   }
 
   void TearDown() override {
-    filesystem_.DeleteDirectoryRecursively(test_dir_.c_str());
+    // Check that the schema store directory is the *only* directory in the
+    // schema_store_dir_. IOW, ensure that all temporary directories have been
+    // properly cleaned up.
+    std::vector<std::string> sub_dirs;
+    ASSERT_TRUE(filesystem_.ListDirectory(temp_dir_.c_str(), &sub_dirs));
+    ASSERT_THAT(sub_dirs, ElementsAre("schema_store"));
+
+    // Finally, clean everything up.
+    ASSERT_TRUE(filesystem_.DeleteDirectoryRecursively(temp_dir_.c_str()));
   }
 
-  const Filesystem filesystem_;
-  const std::string test_dir_;
+  Filesystem filesystem_;
+  std::string temp_dir_;
+  std::string schema_store_dir_;
   SchemaProto schema_;
-  const FakeClock fake_clock_;
+  FakeClock fake_clock_;
 };
 
 TEST_F(SchemaStoreTest, CreationWithNullPointerShouldFail) {
   EXPECT_THAT(
-      SchemaStore::Create(/*filesystem=*/nullptr, test_dir_, &fake_clock_),
+      SchemaStore::Create(/*filesystem=*/nullptr, schema_store_dir_, &fake_clock_),
       StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
 }
 
+TEST_F(SchemaStoreTest, SchemaStoreMoveConstructible) {
+  // Create an instance of SchemaStore.
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("TypeA").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("prop1")
+                  .SetDataTypeString(MATCH_EXACT, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+
+  ICING_ASSERT_OK(schema_store->SetSchema(schema));
+  ICING_ASSERT_OK_AND_ASSIGN(Crc32 expected_checksum,
+                             schema_store->ComputeChecksum());
+
+  // Move construct an instance of SchemaStore
+  SchemaStore move_constructed_schema_store(std::move(*schema_store));
+  EXPECT_THAT(move_constructed_schema_store.GetSchema(),
+              IsOkAndHolds(Pointee(EqualsProto(schema))));
+  EXPECT_THAT(move_constructed_schema_store.ComputeChecksum(),
+              IsOkAndHolds(Eq(expected_checksum)));
+  SectionMetadata expected_metadata(/*id_in=*/0, MATCH_EXACT, TOKENIZER_PLAIN,
+                                    "prop1");
+  EXPECT_THAT(move_constructed_schema_store.GetSectionMetadata("TypeA"),
+              IsOkAndHolds(Pointee(ElementsAre(expected_metadata))));
+}
+
+TEST_F(SchemaStoreTest, SchemaStoreMoveAssignment) {
+  // Create an instance of SchemaStore.
+  SchemaProto schema1 =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("TypeA").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("prop1")
+                  .SetDataTypeString(MATCH_EXACT, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+
+  ICING_ASSERT_OK(schema_store->SetSchema(schema1));
+  ICING_ASSERT_OK_AND_ASSIGN(Crc32 expected_checksum,
+                             schema_store->ComputeChecksum());
+
+  // Construct another instance of SchemaStore
+  SchemaProto schema2 =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("TypeB").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("prop2")
+                  .SetDataTypeString(MATCH_EXACT, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> move_assigned_schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+  ICING_ASSERT_OK(schema_store->SetSchema(schema2));
+
+  // Move assign the first instance into the second one.
+  *move_assigned_schema_store = std::move(*schema_store);
+  EXPECT_THAT(move_assigned_schema_store->GetSchema(),
+              IsOkAndHolds(Pointee(EqualsProto(schema1))));
+  EXPECT_THAT(move_assigned_schema_store->ComputeChecksum(),
+              IsOkAndHolds(Eq(expected_checksum)));
+  SectionMetadata expected_metadata(/*id_in=*/0, MATCH_EXACT, TOKENIZER_PLAIN,
+                                    "prop1");
+  EXPECT_THAT(move_assigned_schema_store->GetSectionMetadata("TypeA"),
+              IsOkAndHolds(Pointee(ElementsAre(expected_metadata))));
+}
+
 TEST_F(SchemaStoreTest, CorruptSchemaError) {
   {
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<SchemaStore> schema_store,
-        SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
     // Set it for the first time
     SchemaStore::SetSchemaResult result;
@@ -121,14 +215,14 @@ TEST_F(SchemaStoreTest, CorruptSchemaError) {
           .AddType(SchemaTypeConfigBuilder().SetType("corrupted"))
           .Build();
 
-  const std::string schema_file = absl_ports::StrCat(test_dir_, "/schema.pb");
+  const std::string schema_file = absl_ports::StrCat(schema_store_dir_, "/schema.pb");
   const std::string serialized_schema = corrupt_schema.SerializeAsString();
 
   filesystem_.Write(schema_file.c_str(), serialized_schema.data(),
                     serialized_schema.size());
 
   // If ground truth was corrupted, we won't know what to do
-  EXPECT_THAT(SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_),
+  EXPECT_THAT(SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_),
               StatusIs(libtextclassifier3::StatusCode::INTERNAL));
 }
 
@@ -136,7 +230,7 @@ TEST_F(SchemaStoreTest, RecoverCorruptDerivedFileOk) {
   {
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<SchemaStore> schema_store,
-        SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
     // Set it for the first time
     SchemaStore::SetSchemaResult result;
@@ -156,12 +250,12 @@ TEST_F(SchemaStoreTest, RecoverCorruptDerivedFileOk) {
   // regenerated from ground truth
 
   const std::string schema_type_mapper_dir =
-      absl_ports::StrCat(test_dir_, "/schema_type_mapper");
+      absl_ports::StrCat(schema_store_dir_, "/schema_type_mapper");
   filesystem_.DeleteDirectoryRecursively(schema_type_mapper_dir.c_str());
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Everything looks fine, ground truth and derived data
   ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_schema,
@@ -174,7 +268,7 @@ TEST_F(SchemaStoreTest, RecoverBadChecksumOk) {
   {
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<SchemaStore> schema_store,
-        SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
     // Set it for the first time
     SchemaStore::SetSchemaResult result;
@@ -193,7 +287,7 @@ TEST_F(SchemaStoreTest, RecoverBadChecksumOk) {
   // the recalculated checksum on initialization. This will force a regeneration
   // of derived files from ground truth.
   const std::string header_file =
-      absl_ports::StrCat(test_dir_, "/schema_store_header");
+      absl_ports::StrCat(schema_store_dir_, "/schema_store_header");
   SchemaStore::Header header;
   header.magic = SchemaStore::Header::kMagic;
   header.checksum = 10;  // Arbitrary garbage checksum
@@ -202,7 +296,7 @@ TEST_F(SchemaStoreTest, RecoverBadChecksumOk) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Everything looks fine, ground truth and derived data
   ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_schema,
@@ -214,7 +308,7 @@ TEST_F(SchemaStoreTest, RecoverBadChecksumOk) {
 TEST_F(SchemaStoreTest, CreateNoPreviousSchemaOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // The apis to retrieve information about the schema should fail gracefully.
   EXPECT_THAT(store->GetSchema(),
@@ -247,7 +341,7 @@ TEST_F(SchemaStoreTest, CreateNoPreviousSchemaOk) {
 TEST_F(SchemaStoreTest, CreateWithPreviousSchemaOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaStore::SetSchemaResult result;
   result.success = true;
@@ -256,7 +350,7 @@ TEST_F(SchemaStoreTest, CreateWithPreviousSchemaOk) {
               IsOkAndHolds(EqualsSetSchemaResult(result)));
 
   schema_store.reset();
-  EXPECT_THAT(SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_),
+  EXPECT_THAT(SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_),
               IsOk());
 }
 
@@ -269,7 +363,7 @@ TEST_F(SchemaStoreTest, MultipleCreateOk) {
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaStore::SetSchemaResult result;
   result.success = true;
@@ -289,7 +383,7 @@ TEST_F(SchemaStoreTest, MultipleCreateOk) {
 
   schema_store.reset();
   ICING_ASSERT_OK_AND_ASSIGN(
-      schema_store, SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      schema_store, SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Verify that our in-memory structures are ok
   EXPECT_THAT(schema_store->GetSchemaTypeConfig("email"),
@@ -305,7 +399,7 @@ TEST_F(SchemaStoreTest, MultipleCreateOk) {
 TEST_F(SchemaStoreTest, SetNewSchemaOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Set it for the first time
   SchemaStore::SetSchemaResult result;
@@ -321,7 +415,7 @@ TEST_F(SchemaStoreTest, SetNewSchemaOk) {
 TEST_F(SchemaStoreTest, SetSameSchemaOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Set it for the first time
   SchemaStore::SetSchemaResult result;
@@ -345,7 +439,7 @@ TEST_F(SchemaStoreTest, SetSameSchemaOk) {
 TEST_F(SchemaStoreTest, SetIncompatibleSchemaOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Set it for the first time
   SchemaStore::SetSchemaResult result;
@@ -372,7 +466,7 @@ TEST_F(SchemaStoreTest, SetIncompatibleSchemaOk) {
 TEST_F(SchemaStoreTest, SetSchemaWithAddedTypeOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto schema = SchemaBuilder()
                            .AddType(SchemaTypeConfigBuilder().SetType("email"))
@@ -406,7 +500,7 @@ TEST_F(SchemaStoreTest, SetSchemaWithAddedTypeOk) {
 TEST_F(SchemaStoreTest, SetSchemaWithDeletedTypeOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto schema =
       SchemaBuilder()
@@ -464,7 +558,7 @@ TEST_F(SchemaStoreTest, SetSchemaWithDeletedTypeOk) {
 TEST_F(SchemaStoreTest, SetSchemaWithReorderedTypesOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto schema =
       SchemaBuilder()
@@ -507,7 +601,7 @@ TEST_F(SchemaStoreTest, SetSchemaWithReorderedTypesOk) {
 TEST_F(SchemaStoreTest, IndexedPropertyChangeRequiresReindexingOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto schema =
       SchemaBuilder()
@@ -551,7 +645,7 @@ TEST_F(SchemaStoreTest, IndexedPropertyChangeRequiresReindexingOk) {
 TEST_F(SchemaStoreTest, IndexNestedDocumentsChangeRequiresReindexingOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Make two schemas. One that sets index_nested_properties to false and one
   // that sets it to true.
@@ -620,7 +714,7 @@ TEST_F(SchemaStoreTest, IndexNestedDocumentsChangeRequiresReindexingOk) {
 TEST_F(SchemaStoreTest, SetSchemaWithIncompatibleTypesOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto schema =
       SchemaBuilder()
@@ -682,7 +776,7 @@ TEST_F(SchemaStoreTest, SetSchemaWithIncompatibleTypesOk) {
 TEST_F(SchemaStoreTest, SetSchemaWithIncompatibleNestedTypesOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // 1. Create a ContactPoint type with a repeated property and set that schema
   SchemaTypeConfigBuilder contact_point_repeated_label =
@@ -749,7 +843,7 @@ TEST_F(SchemaStoreTest, SetSchemaWithIncompatibleNestedTypesOk) {
 TEST_F(SchemaStoreTest, SetSchemaWithIndexIncompatibleNestedTypesOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // 1. Create a ContactPoint type with label that matches prefix and set that
   // schema
@@ -804,7 +898,7 @@ TEST_F(SchemaStoreTest, SetSchemaWithIndexIncompatibleNestedTypesOk) {
 TEST_F(SchemaStoreTest, SetSchemaWithCompatibleNestedTypesOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // 1. Create a ContactPoint type with a optional property and set that schema
   SchemaTypeConfigBuilder contact_point_optional_label =
@@ -857,7 +951,7 @@ TEST_F(SchemaStoreTest, SetSchemaWithCompatibleNestedTypesOk) {
 TEST_F(SchemaStoreTest, GetSchemaTypeId) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   schema_.clear_types();
 
@@ -885,7 +979,7 @@ TEST_F(SchemaStoreTest, GetSchemaTypeId) {
 TEST_F(SchemaStoreTest, ComputeChecksumDefaultOnEmptySchemaStore) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   Crc32 default_checksum;
   EXPECT_THAT(schema_store->ComputeChecksum(), IsOkAndHolds(default_checksum));
@@ -894,7 +988,7 @@ TEST_F(SchemaStoreTest, ComputeChecksumDefaultOnEmptySchemaStore) {
 TEST_F(SchemaStoreTest, ComputeChecksumSameBetweenCalls) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto foo_schema =
       SchemaBuilder().AddType(SchemaTypeConfigBuilder().SetType("foo")).Build();
@@ -910,7 +1004,7 @@ TEST_F(SchemaStoreTest, ComputeChecksumSameBetweenCalls) {
 TEST_F(SchemaStoreTest, ComputeChecksumSameAcrossInstances) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto foo_schema =
       SchemaBuilder().AddType(SchemaTypeConfigBuilder().SetType("foo")).Build();
@@ -923,14 +1017,14 @@ TEST_F(SchemaStoreTest, ComputeChecksumSameAcrossInstances) {
   schema_store.reset();
 
   ICING_ASSERT_OK_AND_ASSIGN(
-      schema_store, SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      schema_store, SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   EXPECT_THAT(schema_store->ComputeChecksum(), IsOkAndHolds(checksum));
 }
 
 TEST_F(SchemaStoreTest, ComputeChecksumChangesOnModification) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto foo_schema =
       SchemaBuilder().AddType(SchemaTypeConfigBuilder().SetType("foo")).Build();
@@ -954,7 +1048,7 @@ TEST_F(SchemaStoreTest, ComputeChecksumChangesOnModification) {
 TEST_F(SchemaStoreTest, PersistToDiskFineForEmptySchemaStore) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Persisting is fine and shouldn't affect anything
   ICING_EXPECT_OK(schema_store->PersistToDisk());
@@ -963,7 +1057,7 @@ TEST_F(SchemaStoreTest, PersistToDiskFineForEmptySchemaStore) {
 TEST_F(SchemaStoreTest, PersistToDiskPreservesAcrossInstances) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   SchemaProto schema =
       SchemaBuilder().AddType(SchemaTypeConfigBuilder().SetType("foo")).Build();
@@ -988,7 +1082,7 @@ TEST_F(SchemaStoreTest, PersistToDiskPreservesAcrossInstances) {
 
   // And we get the same schema back on reinitialization
   ICING_ASSERT_OK_AND_ASSIGN(
-      schema_store, SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      schema_store, SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
   ICING_ASSERT_OK_AND_ASSIGN(actual_schema, schema_store->GetSchema());
   EXPECT_THAT(*actual_schema, EqualsProto(schema));
 }
@@ -996,7 +1090,7 @@ TEST_F(SchemaStoreTest, PersistToDiskPreservesAcrossInstances) {
 TEST_F(SchemaStoreTest, SchemaStoreStorageInfoProto) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Create a schema with two types: one simple type and one type that uses all
   // 16 sections.
@@ -1048,7 +1142,7 @@ TEST_F(SchemaStoreTest, SchemaStoreStorageInfoProto) {
 TEST_F(SchemaStoreTest, GetDebugInfo) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Set schema
   ASSERT_THAT(
@@ -1067,7 +1161,7 @@ TEST_F(SchemaStoreTest, GetDebugInfo) {
 TEST_F(SchemaStoreTest, GetDebugInfoForEmptySchemaStore) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
-      SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
   // Check debug info before setting a schema
   ICING_ASSERT_OK_AND_ASSIGN(SchemaDebugInfoProto out,
@@ -1077,6 +1171,82 @@ TEST_F(SchemaStoreTest, GetDebugInfoForEmptySchemaStore) {
   EXPECT_THAT(out, EqualsProto(expected_out));
 }
 
+TEST_F(SchemaStoreTest, InitializeRegenerateDerivedFilesFailure) {
+  // This test covers the first point that RegenerateDerivedFiles could fail.
+  // This should simply result in SetSchema::Create returning an INTERNAL error.
+
+  {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    SchemaProto schema = SchemaBuilder()
+                             .AddType(SchemaTypeConfigBuilder().SetType("Type"))
+                             .Build();
+    ICING_ASSERT_OK(schema_store->SetSchema(std::move(schema)));
+  }
+
+  auto mock_filesystem = std::make_unique<MockFilesystem>();
+  ON_CALL(*mock_filesystem,
+          CreateDirectoryRecursively(HasSubstr("key_mapper_dir")))
+      .WillByDefault(Return(false));
+  {
+    EXPECT_THAT(SchemaStore::Create(mock_filesystem.get(), schema_store_dir_,
+                                    &fake_clock_),
+                StatusIs(libtextclassifier3::StatusCode::INTERNAL));
+  }
+}
+
+TEST_F(SchemaStoreTest, SetSchemaRegenerateDerivedFilesFailure) {
+  // This test covers the second point that RegenerateDerivedFiles could fail.
+  // If handled correctly, the schema store and section manager should still be
+  // in the original, valid state.
+  SchemaTypeConfigProto type =
+      SchemaTypeConfigBuilder()
+          .SetType("Type")
+          .AddProperty(PropertyConfigBuilder()
+                           .SetName("prop1")
+                           .SetDataTypeString(MATCH_EXACT, TOKENIZER_PLAIN)
+                           .SetCardinality(CARDINALITY_OPTIONAL))
+          .Build();
+  {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
+    SchemaProto schema = SchemaBuilder().AddType(type).Build();
+    ICING_ASSERT_OK(schema_store->SetSchema(std::move(schema)));
+  }
+
+  {
+    auto mock_filesystem = std::make_unique<MockFilesystem>();
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<SchemaStore> schema_store,
+        SchemaStore::Create(mock_filesystem.get(), schema_store_dir_,
+                            &fake_clock_));
+
+    ON_CALL(*mock_filesystem,
+            CreateDirectoryRecursively(HasSubstr("key_mapper_dir")))
+        .WillByDefault(Return(false));
+    SchemaProto schema =
+        SchemaBuilder()
+            .AddType(type)
+            .AddType(SchemaTypeConfigBuilder().SetType("Type2"))
+            .Build();
+    EXPECT_THAT(schema_store->SetSchema(std::move(schema)),
+                StatusIs(libtextclassifier3::StatusCode::INTERNAL));
+    DocumentProto document = DocumentBuilder()
+                                 .SetSchema("Type")
+                                 .AddStringProperty("prop1", "foo bar baz")
+                                 .Build();
+    SectionMetadata expected_metadata(/*id_in=*/0, MATCH_EXACT, TOKENIZER_PLAIN,
+                                      "prop1");
+    ICING_ASSERT_OK_AND_ASSIGN(std::vector<Section> sections,
+                               schema_store->ExtractSections(document));
+    ASSERT_THAT(sections, SizeIs(1));
+    EXPECT_THAT(sections.at(0).metadata, Eq(expected_metadata));
+    EXPECT_THAT(sections.at(0).content, ElementsAre("foo bar baz"));
+  }
+}
+
 }  // namespace
 
 }  // namespace lib
diff --git a/icing/schema/section.h b/icing/schema/section.h
index 40e623a..8b2ba55 100644
--- a/icing/schema/section.h
+++ b/icing/schema/section.h
@@ -77,6 +77,11 @@ struct SectionMetadata {
         id(id_in),
         tokenizer(tokenizer),
         term_match_type(term_match_type_in) {}
+
+  bool operator==(const SectionMetadata& rhs) const {
+    return path == rhs.path && id == rhs.id && tokenizer == rhs.tokenizer &&
+           term_match_type == rhs.term_match_type;
+  }
 };
 
 // Section is an icing internal concept similar to document property but with
diff --git a/icing/scoring/scoring-processor_test.cc b/icing/scoring/scoring-processor_test.cc
index f169039..b42ba31 100644
--- a/icing/scoring/scoring-processor_test.cc
+++ b/icing/scoring/scoring-processor_test.cc
@@ -60,7 +60,7 @@ class ScoringProcessorTest : public testing::Test {
 
     ICING_ASSERT_OK_AND_ASSIGN(
         schema_store_,
-        SchemaStore::Create(&filesystem_, test_dir_, &fake_clock_));
+        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
diff --git a/icing/util/document-validator_test.cc b/icing/util/document-validator_test.cc
index 2261c37..45c23e0 100644
--- a/icing/util/document-validator_test.cc
+++ b/icing/util/document-validator_test.cc
@@ -93,9 +93,11 @@ class DocumentValidatorTest : public ::testing::Test {
                             .SetCardinality(CARDINALITY_REPEATED)))
             .Build();
 
+    schema_dir_ = GetTestTempDir() + "/schema_store";
+    ASSERT_TRUE(filesystem_.CreateDirectory(schema_dir_.c_str()));
     ICING_ASSERT_OK_AND_ASSIGN(
         schema_store_,
-        SchemaStore::Create(&filesystem_, GetTestTempDir(), &fake_clock_));
+        SchemaStore::Create(&filesystem_, schema_dir_, &fake_clock_));
     ASSERT_THAT(schema_store_->SetSchema(schema), IsOk());
 
     document_validator_ =
@@ -122,6 +124,7 @@ class DocumentValidatorTest : public ::testing::Test {
                              SimpleEmailBuilder().Build());
   }
 
+  std::string schema_dir_;
   std::unique_ptr<DocumentValidator> document_validator_;
   std::unique_ptr<SchemaStore> schema_store_;
   Filesystem filesystem_;