From b1273546a555d693535aeabebd158f6925e2961e Mon Sep 17 00:00:00 2001
From: Brandon Roberts <robertsbt@gmail.com>
Date: Sun, 28 Jan 2018 22:27:48 -0600
Subject: [PATCH] fix(Schematics): Add store import to container blueprint

Closes #760
---
 docs/schematics/container.md                   | 4 ++--
 modules/schematics/src/container/index.spec.ts | 9 +++++++++
 modules/schematics/src/container/index.ts      | 8 +++++++-
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/docs/schematics/container.md b/docs/schematics/container.md
index fe66fe1ad4..ed81886c9c 100644
--- a/docs/schematics/container.md
+++ b/docs/schematics/container.md
@@ -25,7 +25,7 @@ ng generate co ComponentName [options]
 
 Provide the path to your file with an exported state interface
 
-- `--reducers`
+- `--state`
   - Type: `string`
 
 Provide the name of the interface exported for your state interface
@@ -39,5 +39,5 @@ Provide the name of the interface exported for your state interface
 Generate a `UsersPage` container component with your reducers imported and the `Store` typed a custom interface named `MyState`.
 
 ```sh
-ng generate container UsersPage --reducers reducers/index.ts --stateInterface MyState
+ng generate container UsersPage --state reducers/index.ts --stateInterface MyState
 ```
diff --git a/modules/schematics/src/container/index.spec.ts b/modules/schematics/src/container/index.spec.ts
index c3601b1135..2a38e3b114 100644
--- a/modules/schematics/src/container/index.spec.ts
+++ b/modules/schematics/src/container/index.spec.ts
@@ -46,6 +46,15 @@ describe('Container Schematic', () => {
     expect(content).toMatch(/import \* as fromStore from '..\/reducers';/);
   });
 
+  it('should import Store into the component', () => {
+    const options = { ...defaultOptions, state: 'reducers' };
+    appTree.create('/src/app/reducers', '');
+    const tree = schematicRunner.runSchematic('container', options, appTree);
+    const content = getFileContent(tree, '/src/app/foo/foo.component.ts');
+    console.log(content);
+    expect(content).toMatch(/import\ {\ Store\ }\ from\ '@ngrx\/store';/);
+  });
+
   it('should update the component constructor if the state path if provided', () => {
     const options = { ...defaultOptions, state: 'reducers' };
     appTree.create('/src/app/reducers', '');
diff --git a/modules/schematics/src/container/index.ts b/modules/schematics/src/container/index.ts
index 9cca957940..19c0700ead 100644
--- a/modules/schematics/src/container/index.ts
+++ b/modules/schematics/src/container/index.ts
@@ -61,6 +61,12 @@ function addStateToComponent(options: FeatureOptions) {
     );
 
     const stateImportPath = buildRelativePath(componentPath, statePath);
+    const storeImport = insertImport(
+      source,
+      componentPath,
+      'Store',
+      '@ngrx/store'
+    );
     const stateImport = options.state
       ? insertImport(
           source,
@@ -94,7 +100,7 @@ function addStateToComponent(options: FeatureOptions) {
       `\n\n  ${storeConstructor}`
     );
 
-    const changes = [stateImport, constructorUpdate];
+    const changes = [storeImport, stateImport, constructorUpdate];
     const recorder = host.beginUpdate(componentPath);
 
     for (const change of changes) {