Пример #1
0
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)
	}
}
Пример #2
0
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)
}