refactor and add a test

backend/remote-state/s3/backend_state.go

@@ -25,28 +25,31 @@ func (b *Backend) States() ([]string, error) {
 		return nil, err
 	}
 
-	envs := []string{backend.DefaultStateName}
+	wss := []string{backend.DefaultStateName}
 	for _, obj := range resp.Contents {
-		env := b.keyEnv(*obj.Key)
-		if env != "" {
-			envs = append(envs, env)
+		ws := getWorkspaceForKey(*obj.Key, b)
+		if ws != "" {
+			wss = append(wss, ws)
 		}
 	}
 
-	sort.Strings(envs[1:])
-	return envs, nil
+	sort.Strings(wss[1:])
+	return wss, nil
 }
 
-// extract the env name from the S3 key
-func (b *Backend) keyEnv(key string) string {
+func getWorkspaceForKey(key string, b *Backend) string {
 	if b.workspaceKeyPrefix == "" {
 		parts := strings.Split(key, "/")
+		if len(parts) > 1 && parts[1] == key {
 			return parts[0]
+		} else {
+			return ""
 		}
+	}
+
 	parts := strings.SplitAfter(key, b.workspaceKeyPrefix)
 
 	if len(parts) < 2 {
-		// no env here
 		return ""
 	}
 

backend/remote-state/s3/backend_test.go

@@ -249,6 +249,62 @@ func TestBackendExtraPaths(t *testing.T) {
 	}
 }
 
+func TestKeyEnv(t *testing.T) {
+	testACC(t)
+	bucketName := fmt.Sprintf("terraform-remote-s3-test-%x", time.Now().Unix())
+	keyName0 := "tfstate"
+	keyName1 := "ws1/tfstate"
+	keyName2 := "ws1/env1/tfstate"
+
+	b0 := backend.TestBackendConfig(t, New(), map[string]interface{}{
+		"bucket":               bucketName,
+		"key":                  keyName0,
+		"encrypt":              true,
+		"workspace_key_prefix": "",
+	}).(*Backend)
+
+	b1 := backend.TestBackendConfig(t, New(), map[string]interface{}{
+		"bucket":               bucketName,
+		"key":                  keyName1,
+		"encrypt":              true,
+		"workspace_key_prefix": "root/userA",
+	}).(*Backend)
+
+	b2 := backend.TestBackendConfig(t, New(), map[string]interface{}{
+		"bucket":               bucketName,
+		"key":                  keyName2,
+		"encrypt":              true,
+		"workspace_key_prefix": "root/userA",
+	}).(*Backend)
+
+	if err := testGetWorkspaceForKey(b0, "tfstate", ""); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := testGetWorkspaceForKey(b0, "ws1/tfstate", "ws1"); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := testGetWorkspaceForKey(b1, "root/userA/ws1/tfstate", "ws1"); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := testGetWorkspaceForKey(b1, "root/userA/ws2/tfstate", "ws2"); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := testGetWorkspaceForKey(b2, "root/userA/ws2/env1/tfstate", "ws2"); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func testGetWorkspaceForKey(b *Backend, key string, expected string) error {
+	if getWorkspaceForKey(key, b) != expected {
+		return fmt.Errorf("incorrect workspace for key[%q]: %q", expected, key)
+	}
+	return nil
+}
+
 func checkStateList(b backend.Backend, expected []string) error {
 	states, err := b.States()
 	if err != nil {