From 9c96a532b1f59af2d3939c26ef8bf8a85044b46b Mon Sep 17 00:00:00 2001
From: Pmyl <julo134@gmail.com>
Date: Sun, 13 Sep 2020 17:40:46 +0100
Subject: [PATCH] fix(enum): fix enum constant computed properties

---
 package.json                                  |  1 +
 .../descriptor/enum/enumDeclaration.ts        | 26 +++++-----
 test/features/random/enum.test.ts             | 47 ++++++++++++++++---
 test/playground/jasmine.build.json            |  5 ++
 4 files changed, 58 insertions(+), 21 deletions(-)
 create mode 100644 test/playground/jasmine.build.json

diff --git a/package.json b/package.json
index d176d3e4c..df01aadd8 100644
--- a/package.json
+++ b/package.json
@@ -22,6 +22,7 @@
         "test:features": "cross-env JASMINE_CONFIG=./test/features/jasmine.json TSCONFIG=./test/features/tsconfig.json npm run test:common",
         "test:logs": "cross-env JASMINE_CONFIG=./test/logs/jasmine.json TSCONFIG=./test/logs/tsconfig.json npm run test:common",
         "test:playground": "cross-env JASMINE_CONFIG=./test/playground/jasmine.json TSCONFIG=./test/playground/tsconfig.json npm run test:common",
+        "test:playground:build": "cross-env JASMINE_CONFIG=./test/playground/jasmine.build.json TSCONFIG=./test/playground/tsconfig.json npm run test:common",
         "test:common": "cross-var ts-node --files -r tsconfig-paths/register --compiler ttypescript --project $TSCONFIG node_modules/jasmine/bin/jasmine --config=$JASMINE_CONFIG",
         "dist:collect": "cp -r package.json package-lock.json README.md dist",
         "ts-check:src": "tsc --noEmit",
diff --git a/src/transformer/descriptor/enum/enumDeclaration.ts b/src/transformer/descriptor/enum/enumDeclaration.ts
index c06499118..9982d023b 100644
--- a/src/transformer/descriptor/enum/enumDeclaration.ts
+++ b/src/transformer/descriptor/enum/enumDeclaration.ts
@@ -7,19 +7,11 @@ export function GetEnumDeclarationDescriptor(
   node: ts.EnumDeclaration
 ): ts.Expression {
   const typeChecker: ts.TypeChecker = TypeChecker();
-  const typesList: ts.LiteralType[] = node.members.map((it: ts.EnumMember) =>
-    typeChecker.getTypeAtLocation(it)
-  ) as ts.LiteralType[];
 
   if (IsTsAutoMockRandomEnabled()) {
-    const nodesList: ts.Expression[] = typesList.map(
-      (type: ts.LiteralType, index: number) => {
-        if (type.hasOwnProperty('value')) {
-          return ts.createLiteral(type.value);
-        }
-
-        return ts.createLiteral(index);
-      }
+    const nodesList: ts.Expression[] = node.members.map(
+      (member: ts.EnumMember, index: number) =>
+        getEnumMemberValue(typeChecker, member, index)
     );
 
     return ts.createCall(
@@ -29,9 +21,13 @@ export function GetEnumDeclarationDescriptor(
     );
   }
 
-  if (typesList[0].hasOwnProperty('value')) {
-    return ts.createLiteral(typesList[0].value);
-  }
+  return getEnumMemberValue(typeChecker, node.members[0]);
+}
 
-  return ts.createLiteral(0);
+function getEnumMemberValue(
+  typeChecker: ts.TypeChecker,
+  member: ts.EnumMember,
+  defaultValue: string | number = 0
+): ts.Expression {
+  return ts.createLiteral(typeChecker.getConstantValue(member) || defaultValue);
 }
diff --git a/test/features/random/enum.test.ts b/test/features/random/enum.test.ts
index c6d80a287..2fb1a01c1 100644
--- a/test/features/random/enum.test.ts
+++ b/test/features/random/enum.test.ts
@@ -96,15 +96,50 @@ describe('Random enum', () => {
     expect(thirdMock.direction).toBe(Direction.None);
   });
 
-  it('should have position of "computed" property, because transformer don\'t support computed Enum values', () => {
+  it('should get the correct const enum member computed property', () => {
+    const enum WithComputed {
+      Computed = 1 + 4,
+    }
+
+    interface IWithComputed {
+      computed: WithComputed;
+    }
+
+    const spy: jasmine.Spy = spyOn(Math, 'floor');
+
+    spy.and.returnValue(0);
+    const mock1: IWithComputed = createMock<IWithComputed>();
+
+    expect(mock1.computed).toBe(5);
+  });
+
+  it('should get the correct enum member computed property for constant expression', () => {
+    enum WithComputed {
+      Computed = 1 + 4,
+    }
+
+    interface IWithComputed {
+      computed: WithComputed;
+    }
+
+    const spy: jasmine.Spy = spyOn(Math, 'floor');
+
+    spy.and.returnValue(0);
+    const mock1: IWithComputed = createMock<IWithComputed>();
+
+    expect(mock1.computed).toBe(5);
+  });
+
+  it('should have position of non constant "computed" property, because transformer don\'t support non constant computed Enum values', () => {
     // Issue https://github.com/Typescript-TDD/ts-auto-mock/issues/339
     function getComputed(): number {
       return 12;
     }
 
     enum WithComputed {
+      NonComputed,
       Computed = getComputed(),
-      Computed2 = 1 + 4,
+      Computed1 = getComputed(),
     }
 
     interface IWithComputed {
@@ -113,14 +148,14 @@ describe('Random enum', () => {
 
     const spy: jasmine.Spy = spyOn(Math, 'floor');
 
-    spy.and.returnValue(0);
+    spy.and.returnValue(1);
     const mock1: IWithComputed = createMock<IWithComputed>();
 
-    expect(mock1.computed).toBe(0);
+    expect(mock1.computed).toBe(1);
 
-    spy.and.returnValue(1);
+    spy.and.returnValue(2);
     const mock2: IWithComputed = createMock<IWithComputed>();
 
-    expect(mock2.computed).toBe(1);
+    expect(mock2.computed).toBe(2);
   });
 });
diff --git a/test/playground/jasmine.build.json b/test/playground/jasmine.build.json
new file mode 100644
index 000000000..540e603cc
--- /dev/null
+++ b/test/playground/jasmine.build.json
@@ -0,0 +1,5 @@
+{
+  "spec_dir": "./test/playground",
+  "spec_files": ["**/*test.js"],
+  "helpers": ["../reporter.js"]
+}