From d66597a41fe26a4eb9aeadd465fc529d4ea66ddd Mon Sep 17 00:00:00 2001
From: libin14 <>
Date: Mon, 12 Oct 2020 19:46:24 +0800
Subject: [PATCH] fix  store cache in memory bug

---
 SDWebImage/Core/SDImageCache.m  | 18 +++++++++++++-----
 Tests/Tests/SDImageCacheTests.m | 12 ++++++++++++
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/SDWebImage/Core/SDImageCache.m b/SDWebImage/Core/SDImageCache.m
index 6ef27dd5..84090f08 100644
--- a/SDWebImage/Core/SDImageCache.m
+++ b/SDWebImage/Core/SDImageCache.m
@@ -523,11 +523,19 @@
                 cacheType = SDImageCacheTypeMemory;
             } else if (diskData) {
                 cacheType = SDImageCacheTypeDisk;
-                // decode image data only if in-memory cache missed
-                diskImage = [self diskImageForKey:key data:diskData options:options context:context];
-                if (diskImage && self.config.shouldCacheImagesInMemory) {
-                    NSUInteger cost = diskImage.sd_memoryCost;
-                    [self.memoryCache setObject:diskImage forKey:key cost:cost];
+                
+                BOOL shouldCacheToMomery = YES;
+                if (context[SDWebImageContextStoreCacheType]) {
+                    SDImageCacheType cacheType = [context[SDWebImageContextStoreCacheType] integerValue];
+                    shouldCacheToMomery = (cacheType == SDImageCacheTypeAll || cacheType == SDImageCacheTypeMemory);
+                }
+                if (shouldCacheToMomery) {
+                    // decode image data only if in-memory cache missed
+                    diskImage = [self diskImageForKey:key data:diskData options:options context:context];
+                    if (diskImage && self.config.shouldCacheImagesInMemory) {
+                        NSUInteger cost = diskImage.sd_memoryCost;
+                        [self.memoryCache setObject:diskImage forKey:key cost:cost];
+                    }
                 }
             }
             
diff --git a/Tests/Tests/SDImageCacheTests.m b/Tests/Tests/SDImageCacheTests.m
index 96fb0f48..f57b12db 100644
--- a/Tests/Tests/SDImageCacheTests.m
+++ b/Tests/Tests/SDImageCacheTests.m
@@ -607,6 +607,18 @@ static NSString *kTestImageKeyPNG = @"TestImageKey.png";
 }
 
 #pragma mark - SDImageCache & SDImageCachesManager
+- (void)test49SDImageCacheQueryOp {
+    XCTestExpectation *expectation = [self expectationWithDescription:@"SDImageCache query op works"];
+    NSData *data = [[SDImageCodersManager sharedManager] encodedDataWithImage:[self testJPEGImage] format:SDImageFormatJPEG options:nil];
+    [[SDImageCache sharedImageCache] storeImageDataToDisk:data forKey:kTestImageKeyJPEG];
+    
+    [[SDImageCachesManager sharedManager] queryImageForKey:kTestImageKeyJPEG options:0 context:@{SDWebImageContextStoreCacheType : @(SDImageCacheTypeDisk)} cacheType:SDImageCacheTypeAll completion:^(UIImage * _Nullable image, NSData * _Nullable data, SDImageCacheType cacheType) {
+        expect([[SDImageCache sharedImageCache] imageFromMemoryCacheForKey:kTestImageKeyJPEG]).beNil();
+        [expectation fulfill];
+    }];
+    [self waitForExpectationsWithCommonTimeout];
+}
+
 - (void)test50SDImageCacheQueryOp {
     XCTestExpectation *expectation = [self expectationWithDescription:@"SDImageCache query op works"];
     [[SDImageCache sharedImageCache] storeImage:[self testJPEGImage] forKey:kTestImageKeyJPEG toDisk:NO completion:nil];