1/*
2 * Copyright (C) 2015 Square, Inc.
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 */
16package com.squareup.okhttp.internal.io;
17
18import java.io.File;
19import java.io.FileNotFoundException;
20import java.io.IOException;
21import java.util.ArrayList;
22import java.util.IdentityHashMap;
23import java.util.Iterator;
24import java.util.LinkedHashMap;
25import java.util.List;
26import java.util.Map;
27import okio.Buffer;
28import okio.ForwardingSink;
29import okio.ForwardingSource;
30import okio.Sink;
31import okio.Source;
32import org.junit.rules.TestRule;
33import org.junit.runner.Description;
34import org.junit.runners.model.Statement;
35
36/** A simple file system where all files are held in memory. Not safe for concurrent use. */
37public final class InMemoryFileSystem implements FileSystem, TestRule {
38  private final Map<File, Buffer> files = new LinkedHashMap<>();
39  private final Map<Source, File> openSources = new IdentityHashMap<>();
40  private final Map<Sink, File> openSinks = new IdentityHashMap<>();
41
42  @Override public Statement apply(final Statement base, Description description) {
43    return new Statement() {
44      @Override public void evaluate() throws Throwable {
45        base.evaluate();
46        ensureResourcesClosed();
47      }
48    };
49  }
50
51  public void ensureResourcesClosed() {
52    List<String> openResources = new ArrayList<>();
53    for (File file : openSources.values()) {
54      openResources.add("Source for " + file);
55    }
56    for (File file : openSinks.values()) {
57      openResources.add("Sink for " + file);
58    }
59    if (!openResources.isEmpty()) {
60      StringBuilder builder = new StringBuilder("Resources acquired but not closed:");
61      for (String resource : openResources) {
62        builder.append("\n * ").append(resource);
63      }
64      throw new IllegalStateException(builder.toString());
65    }
66  }
67
68  @Override public Source source(File file) throws FileNotFoundException {
69    Buffer result = files.get(file);
70    if (result == null) throw new FileNotFoundException();
71
72    final Source source = result.clone();
73    openSources.put(source, file);
74
75    return new ForwardingSource(source) {
76      @Override public void close() throws IOException {
77        openSources.remove(source);
78        super.close();
79      }
80    };
81  }
82
83  @Override public Sink sink(File file) throws FileNotFoundException {
84    return sink(file, false);
85  }
86
87  @Override public Sink appendingSink(File file) throws FileNotFoundException {
88    return sink(file, true);
89  }
90
91  private Sink sink(File file, boolean appending) {
92    Buffer result = null;
93    if (appending) {
94      result = files.get(file);
95    }
96    if (result == null) {
97      result = new Buffer();
98    }
99    files.put(file, result);
100
101    final Sink sink = result;
102    openSinks.put(sink, file);
103
104    return new ForwardingSink(sink) {
105      @Override public void close() throws IOException {
106        openSinks.remove(sink);
107        super.close();
108      }
109    };
110  }
111
112  @Override public void delete(File file) throws IOException {
113    files.remove(file);
114  }
115
116  @Override public boolean exists(File file) throws IOException {
117    return files.containsKey(file);
118  }
119
120  @Override public long size(File file) {
121    Buffer buffer = files.get(file);
122    return buffer != null ? buffer.size() : 0L;
123  }
124
125  @Override public void rename(File from, File to) throws IOException {
126    Buffer buffer = files.remove(from);
127    if (buffer == null) throw new FileNotFoundException();
128    files.put(to, buffer);
129  }
130
131  @Override public void deleteContents(File directory) throws IOException {
132    String prefix = directory.toString() + "/";
133    for (Iterator<File> i = files.keySet().iterator(); i.hasNext(); ) {
134      File file = i.next();
135      if (file.toString().startsWith(prefix)) i.remove();
136    }
137  }
138}
139