package org.apidesign.exceptions.trycatchredo;

 import java.io.ByteArrayInputStream;

 import java.io.ByteArrayOutputStream;

 import java.io.IOException;

 import java.io.InputStream;

 import java.io.OutputStream;

 import java.net.MalformedURLException;

 import java.net.URL;

 import java.net.URLConnection;

 import java.net.URLStreamHandler;

 import java.net.URLStreamHandlerFactory;

 import java.util.HashMap;

 import java.util.Map;

 import junit.framework.Assert;

 /**

  *

  * @author Jaroslav Tulach

  */

 public final class MemoryURL extends URLStreamHandler {

     private MemoryURL() {

     }

     static void initialize() {

     }

     static {

         class F implements URLStreamHandlerFactory {

             public URLStreamHandler createURLStreamHandler(String protocol) {

                 if (protocol.startsWith("memory")) {

                     return new MemoryURL();

                 }

                 return null;

             }

         }

         F f = new F();

         URL.setURLStreamHandlerFactory(f);

     }

     private static Map<String,InputStream> contents = new HashMap<String,InputStream>();

     private static Map<String,MC> outputs = new HashMap<String,MC>();

     public static void registerURL(String u, String content, ByteArrayOutputStream out) throws MalformedURLException {

         contents.put(u, new ByteArrayInputStream(content.getBytes()));

         if (out != null) {

             new MC(new URL(u)).out = out;

         }

     }

     public static byte[] getOutputForURL(String u) {

         MC out = outputs.get(u);

         Assert.assertNotNull("No output for " + u, out);

         return out.out.toByteArray();

     }

     protected URLConnection openConnection(URL u) throws IOException {

         return new MC(u);

     }

     private static final class MC extends URLConnection {

         private InputStream values;

         private ByteArrayOutputStream out;

         public MC(URL u) {

             super(u);

             outputs.put(u.toExternalForm(), this);

         }

         public void connect() throws IOException {

             if (values != null) {

                 return;

             }

             values = contents.remove(url.toExternalForm());

             if (values == null) {

                 throw new IOException("No such content: " + url);

             }

         }

         @Override

         public InputStream getInputStream() throws IOException {

             connect();

             return values;

         }

         @Override

         public OutputStream getOutputStream() throws IOException {

             if (out == null) {

                 out = new ByteArrayOutputStream();

             }

             return out;

         }

     }

 }