func main() { flag.Usage = func() { fmt.Fprintf(os.Stderr, usage, os.Args[0]) os.Exit(2) } flag.Parse() if *printVersion { fmt.Println(version) os.Exit(0) } appLogger, err := setLogger(os.Getenv("AMIQUERY_APP_LOGFILE")) if err != nil { log.Fatalln("Unable to set application logging output:", err) } log.SetOutput(appLogger) httpLogger, err := setLogger(os.Getenv("AMIQUERY_HTTP_LOGFILE")) if err != nil { log.Fatalln("Unable to set HTTP logging output:", err) } cfg, err := NewConfig() if err != nil { log.Fatal(err) } log.Println("Loaded configuration") cacheMgr, err := amicache.NewManager(cfg.Manager, amicache.TTL(cfg.CacheTTL), amicache.Regions(cfg.Regions...), amicache.OwnerIDs(cfg.OwnerIDs...), amicache.AssumeRole(cfg.RoleARN), amicache.HTTPClient(&http.Client{ Transport: &http.Transport{ Proxy: http.ProxyFromEnvironment, }, }), ) if err != nil { log.Fatal(err) } if *daemonize { log.Println("Daemonizing process...") godaemon.MakeDaemon(&godaemon.DaemonAttr{}) } cacheMgr.Start() defer cacheMgr.Stop() // Create a context and add the cache manager ctx := context.Background() ctx = context.WithValue(ctx, api.CacheManagerKey, cacheMgr) // Version 1 of the REST API v1 := &api.ContextAdapter{ Context: ctx, Handler: api.ContextHandlerFunc(v1.Handler), } router := mux.NewRouter() // Add the version 1 handler and set it as the default (Accept: */*) router.Handle("/amis", handlers.CombinedLoggingHandler(httpLogger, v1)). HeadersRegexp("Accept", `(application/vnd\.ami-query-v1\+json|\*/\*)`). Methods("GET") http.Handle("/", router) if err := http.ListenAndServe(cfg.ListenAddr, nil); err != nil { log.Fatal(err) } }
func TestMain(m *testing.M) { // Load test data usEast1, err := ioutil.ReadFile("../../testdata/us-east-1-describe-images.xml") if err != nil { panic(err) } usWest1, err := ioutil.ReadFile("../../testdata/us-west-1-describe-images.xml") if err != nil { panic(err) } usWest2, err := ioutil.ReadFile("../../testdata/us-west-2-describe-images.xml") if err != nil { panic(err) } // Mock EC2 service endpoint mux := http.NewServeMux() server := httptest.NewServer(mux) // Set EC2Endpoint to the mocked endpoint and reset it afterwards. defer func(previous func(string) string) { amicache.EC2Endpoint = previous }(amicache.EC2Endpoint) amicache.EC2Endpoint = func(region string) string { return fmt.Sprintf("%s/%s", server.URL, region) } // Endpoint response handlers mux.HandleFunc("/us-east-1/", func(w http.ResponseWriter, r *http.Request) { fmt.Fprintf(w, "%s", usEast1) }) mux.HandleFunc("/us-west-1/", func(w http.ResponseWriter, r *http.Request) { fmt.Fprintf(w, "%s", usWest1) }) mux.HandleFunc("/us-west-2/", func(w http.ResponseWriter, r *http.Request) { fmt.Fprintf(w, "%s", usWest2) }) creds := credentials.NewStaticCredentials("foo", "bar", "baz") // Fill cache with test data cacheMgr, _ := amicache.NewManager( amicache.NewInternalManager(), amicache.AWSCreds(creds), amicache.Regions("us-east-1", "us-west-1", "us-west-2"), ) <-cacheMgr.Start() // Mock amiquery service amiquery = httptest.NewServer(&api.ContextAdapter{ Context: context.WithValue(context.Background(), api.CacheManagerKey, cacheMgr), Handler: api.ContextHandlerFunc(Handler), }) // Run tests rc := m.Run() // Shutdown services amiquery.Close() server.Close() cacheMgr.Stop() os.Exit(rc) }