Example #1
0
// NewConfig returns a Config with settings pulled from the environment. See
// the README.md for more information.
func NewConfig() (*Config, error) {
	var err error
	var cfg = &Config{
		ListenAddr: ":8080",
		CacheTTL:   15 * time.Minute,
	}

	// Set the listen address
	if laddr := os.Getenv("AMIQUERY_LISTEN_ADDRESS"); laddr != "" {
		cfg.ListenAddr = laddr
	}

	// AWS regions to scan for AMIs
	if regions := os.Getenv("AMIQUERY_REGIONS"); regions != "" {
		cfg.Regions = strings.Split(regions, ",")
	}

	// Duration between cache updates
	if ttl := os.Getenv("AMIQUERY_CACHE_TTL"); ttl != "" {
		if cfg.CacheTTL, err = time.ParseDuration(ttl); err != nil {
			return nil, err
		}
	}

	// Owner IDs used to filter AMI results
	if ownerIDs := os.Getenv("AMIQUERY_OWNER_IDS"); ownerIDs != "" {
		cfg.OwnerIDs = strings.Split(ownerIDs, ",")
	} else {
		return nil, fmt.Errorf("AMIQUERY_OWNER_IDS is undefined")
	}

	// ARN used for AssumeRole
	if arn := os.Getenv("AMIQUERY_ROLE_ARN"); arn != "" {
		cfg.RoleARN = arn
	}

	// The type of underlying cache to use
	cacheType := os.Getenv("AMIQUERY_CACHE_MANAGER")
	switch cacheType {
	case "memcached":
		servers := strings.Split(os.Getenv("AMIQUERY_MEMCACHED_SERVERS"), ",")
		if len(servers) == 0 {
			return nil, fmt.Errorf("AMIQUERY_MEMCACHED_SERVERS is undefined")
		}
		cfg.Manager = amicache.NewMemcachedManager(servers...)
	default:
		cfg.Manager = amicache.NewInternalManager()
	}

	return cfg, nil
}
Example #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)
}