1 /* 2 * Copyright (C) 2017 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #include "zlib-utils.h" 18 19 #include <memory> 20 21 #include "model_generated.h" 22 #include "gmock/gmock.h" 23 #include "gtest/gtest.h" 24 25 namespace libtextclassifier2 { 26 27 TEST(ZlibUtilsTest, CompressModel) { 28 ModelT model; 29 model.regex_model.reset(new RegexModelT); 30 model.regex_model->patterns.emplace_back(new RegexModel_::PatternT); 31 model.regex_model->patterns.back()->pattern = "this is a test pattern"; 32 model.regex_model->patterns.emplace_back(new RegexModel_::PatternT); 33 model.regex_model->patterns.back()->pattern = "this is a second test pattern"; 34 35 model.datetime_model.reset(new DatetimeModelT); 36 model.datetime_model->patterns.emplace_back(new DatetimeModelPatternT); 37 model.datetime_model->patterns.back()->regexes.emplace_back( 38 new DatetimeModelPattern_::RegexT); 39 model.datetime_model->patterns.back()->regexes.back()->pattern = 40 "an example datetime pattern"; 41 model.datetime_model->extractors.emplace_back(new DatetimeModelExtractorT); 42 model.datetime_model->extractors.back()->pattern = 43 "an example datetime extractor"; 44 45 // Compress the model. 46 EXPECT_TRUE(CompressModel(&model)); 47 48 // Sanity check that uncompressed field is removed. 49 EXPECT_TRUE(model.regex_model->patterns[0]->pattern.empty()); 50 EXPECT_TRUE(model.regex_model->patterns[1]->pattern.empty()); 51 EXPECT_TRUE(model.datetime_model->patterns[0]->regexes[0]->pattern.empty()); 52 EXPECT_TRUE(model.datetime_model->extractors[0]->pattern.empty()); 53 54 // Pack and load the model. 55 flatbuffers::FlatBufferBuilder builder; 56 builder.Finish(Model::Pack(builder, &model)); 57 const Model* compressed_model = 58 GetModel(reinterpret_cast<const char*>(builder.GetBufferPointer())); 59 ASSERT_TRUE(compressed_model != nullptr); 60 61 // Decompress the fields again and check that they match the original. 62 std::unique_ptr<ZlibDecompressor> decompressor = ZlibDecompressor::Instance(); 63 ASSERT_TRUE(decompressor != nullptr); 64 std::string uncompressed_pattern; 65 EXPECT_TRUE(decompressor->Decompress( 66 compressed_model->regex_model()->patterns()->Get(0)->compressed_pattern(), 67 &uncompressed_pattern)); 68 EXPECT_EQ(uncompressed_pattern, "this is a test pattern"); 69 EXPECT_TRUE(decompressor->Decompress( 70 compressed_model->regex_model()->patterns()->Get(1)->compressed_pattern(), 71 &uncompressed_pattern)); 72 EXPECT_EQ(uncompressed_pattern, "this is a second test pattern"); 73 EXPECT_TRUE(decompressor->Decompress(compressed_model->datetime_model() 74 ->patterns() 75 ->Get(0) 76 ->regexes() 77 ->Get(0) 78 ->compressed_pattern(), 79 &uncompressed_pattern)); 80 EXPECT_EQ(uncompressed_pattern, "an example datetime pattern"); 81 EXPECT_TRUE(decompressor->Decompress(compressed_model->datetime_model() 82 ->extractors() 83 ->Get(0) 84 ->compressed_pattern(), 85 &uncompressed_pattern)); 86 EXPECT_EQ(uncompressed_pattern, "an example datetime extractor"); 87 88 EXPECT_TRUE(DecompressModel(&model)); 89 EXPECT_EQ(model.regex_model->patterns[0]->pattern, "this is a test pattern"); 90 EXPECT_EQ(model.regex_model->patterns[1]->pattern, 91 "this is a second test pattern"); 92 EXPECT_EQ(model.datetime_model->patterns[0]->regexes[0]->pattern, 93 "an example datetime pattern"); 94 EXPECT_EQ(model.datetime_model->extractors[0]->pattern, 95 "an example datetime extractor"); 96 } 97 98 } // namespace libtextclassifier2 99
