From 40b870576a4fc339380ddb81357a513fa7ce6fc5 Mon Sep 17 00:00:00 2001
From: Ivan Pomykacz <ivan.pomykacz@lmc.eu>
Date: Thu, 21 Mar 2024 15:22:46 +0200
Subject: [PATCH] chore: resolves #14

---
 src/project/project-context.ts |  9 ++++
 test/helpers/app.ts            |  8 +++-
 test/wrapper.test.ts           | 83 ++++++++++++++++++++++++++++++++++
 3 files changed, 99 insertions(+), 1 deletion(-)
 create mode 100644 test/wrapper.test.ts

diff --git a/src/project/project-context.ts b/src/project/project-context.ts
index f7f7ace..bb2e869 100644
--- a/src/project/project-context.ts
+++ b/src/project/project-context.ts
@@ -5,6 +5,7 @@ import { Account } from './interfaces';
 import { Project } from './project';
 import { addError } from '../error';
 import { get } from '../utils/get';
+import { isNonEmptyString } from '../utils/isNonEmptyString';
 
 export class ProjectContext {
   /**
@@ -66,6 +67,14 @@ export class ProjectContext {
   }
 
   static getEnvironment(scope: Construct): string {
+    const environmentType = ProjectContext.tryGetEnvironment(scope);
+    const tryAccountType: string = ProjectContext.getAccountType(scope);
+    if (isNonEmptyString(environmentType) && tryAccountType == '') {
+      const accountType = ProjectContext._getAccountTypeByEnvironment(scope, environmentType);
+      if (isNonEmptyString(accountType)) {
+        AccountType.set(scope, accountType);
+      }
+    }
     const allowedEnvironments = ProjectContext.getAllowedEnvironments(scope);
     return EnvironmentType.get(scope, allowedEnvironments);
   }
diff --git a/test/helpers/app.ts b/test/helpers/app.ts
index 86cf03b..7c5f604 100644
--- a/test/helpers/app.ts
+++ b/test/helpers/app.ts
@@ -6,7 +6,7 @@ import * as route53 from 'aws-cdk-lib/aws-route53';
 import * as s3 from 'aws-cdk-lib/aws-s3';
 import * as ssm from 'aws-cdk-lib/aws-ssm';
 import { Construct } from 'constructs';
-import { Project, EnvironmentWrapper, SmartStack, EC, Name, PathName, UrlName, ProjectProps, AC } from '../../src';
+import { Project, EnvironmentWrapper, SmartStack, EC, Name, PathName, UrlName, ProjectProps, AC, AccountWrapper } from '../../src';
 
 
 export class MyStack extends SmartStack {
@@ -52,6 +52,12 @@ export class MyStack extends SmartStack {
   }
 }
 
+export class Account extends AccountWrapper {
+  constructor(scope: Construct) {
+    super(scope);
+  }
+}
+
 export class Environment extends EnvironmentWrapper {
   constructor(scope: Construct) {
     super(scope);
diff --git a/test/wrapper.test.ts b/test/wrapper.test.ts
new file mode 100644
index 0000000..aa0d9e5
--- /dev/null
+++ b/test/wrapper.test.ts
@@ -0,0 +1,83 @@
+import { Account, Environment } from './helpers/app';
+import { AccountStrategy, PC, Project } from '../src';
+
+
+describe('Account and Environment wrappers', () => {
+
+  const props = {
+    name: 'my-cool-project',
+    author: {
+      organization: 'Acme Corp',
+      name: 'Mad Scientists',
+      email: 'mad.scientists@acme.example.com',
+    },
+    defaultRegion: 'eu-west-1',
+    accounts: AccountStrategy.two({
+      dev: {
+        id: '111111111111',
+      },
+      prod: {
+        id: '222222222222',
+      },
+    }),
+  };
+
+  describe('AccountWrapper', () => {
+
+    test('dev/account', () => {
+      const project = new Project(props);
+      project.node.setContext('account-type', 'dev');
+      const account = new Account(project);
+      expect(PC.getAccountType(account)).toBe('dev');
+    });
+
+    test('prod/account', () => {
+      const project = new Project(props);
+      project.node.setContext('account-type', 'prod');
+      const account = new Account(project);
+      expect(PC.getAccountType(account)).toBe('prod');
+    });
+
+  });
+
+  describe('EnvironmentWrapper', () => {
+
+    ['development', 'feature/foo-bar', 'test', 'staging'].forEach((env) => {
+      test('dev/environment/with-account', () => {
+        const project = new Project(props);
+        project.node.setContext('account-type', 'dev');
+        project.node.setContext('environment-type', env);
+        const environment = new Environment(project);
+        expect(PC.getAccountType(environment)).toBe('dev');
+        expect(PC.getEnvironment(environment)).toBe(env);
+      });
+
+      test('dev/environment/no-account', () => {
+        const project = new Project(props);
+        project.node.setContext('environment-type', env);
+        const environment = new Environment(project);
+        expect(PC.getAccountType(environment)).toBe('dev');
+        expect(PC.getEnvironment(environment)).toBe(env);
+      });
+    });
+
+    ['preproduction', 'production'].forEach((env) => {
+      test('prod/environment/with-account', () => {
+        const project = new Project(props);
+        project.node.setContext('account-type', 'prod');
+        project.node.setContext('environment-type', env);
+        const environment = new Environment(project);
+        expect(PC.getAccountType(environment)).toBe('prod');
+        expect(PC.getEnvironment(environment)).toBe(env);
+      });
+
+      test('prod/environment/no-account', () => {
+        const project = new Project(props);
+        project.node.setContext('environment-type', env);
+        const environment = new Environment(project);
+        expect(PC.getAccountType(environment)).toBe('prod');
+        expect(PC.getEnvironment(environment)).toBe(env);
+      });
+    });
+  });
+});