From b922750de07ab5c2865c599288aa591560c46f5c Mon Sep 17 00:00:00 2001
From: beelit94 <beelit94@gmail.com>
Date: Tue, 9 May 2017 17:47:36 +0800
Subject: [PATCH] use log to test

---
 test/import_test/test.tf        |  7 +++++
 test/test_terraform.py          | 50 ++++++++++++++++++++++++---------
 test/vars_require_input/main.tf |  4 +--
 3 files changed, 45 insertions(+), 16 deletions(-)
 create mode 100644 test/import_test/test.tf

diff --git a/test/import_test/test.tf b/test/import_test/test.tf
new file mode 100644
index 0000000..9da2429
--- /dev/null
+++ b/test/import_test/test.tf
@@ -0,0 +1,7 @@
+provider "archive" {}
+
+data "archive_file" "init" {
+  type        = "zip"
+  source_file = "${path.module}/init.tpl"
+  output_path = "${path.module}/init.zip"
+}
diff --git a/test/test_terraform.py b/test/test_terraform.py
index 727fc80..1a1dba5 100644
--- a/test/test_terraform.py
+++ b/test/test_terraform.py
@@ -1,3 +1,7 @@
+try:
+    from cStringIO import StringIO      # Python 2
+except ImportError:
+    from io import StringIO
 from python_terraform import *
 import pytest
 import os
@@ -6,8 +10,9 @@ import re
 import shutil
 
 logging.basicConfig(level=logging.DEBUG)
-ch = logging.StreamHandler(sys.stdout)
-logging.getLogger().addHandler(ch)
+root_logger = logging.getLogger()
+# ch = logging.StreamHandler(sys.stdout)
+# root_logger.addHandler(ch)
 current_path = os.path.dirname(os.path.realpath(__file__))
 
 STRING_CASES = [
@@ -25,18 +30,20 @@ STRING_CASES = [
  ]
 
 CMD_CASES = [
-    ['method', 'expected_output', 'expected_ret_code'],
+    ['method', 'expected_output', 'expected_ret_code', 'expected_logs'],
     [
         [
             lambda x: x.cmd('plan', 'var_to_output', no_color=IsFlagged, var={'test_var': 'test'}) ,
             "doesn't need to do anything",
-            0
+            0,
+            ''
         ],
         # try import aws instance
         [
             lambda x: x.cmd('import', 'aws_instance.foo', 'i-abcd1234', no_color=IsFlagged),
-            'Import complete!',
-            1
+            '',
+            1,
+            'command: terraform import -no-color aws_instance.foo i-abcd1234'
         ]
     ]
 ]
@@ -56,6 +63,19 @@ def fmt_test_file(request):
     return
 
 
+@pytest.fixture()
+def string_logger(request):
+    log_stream = StringIO()
+    handler = logging.StreamHandler(log_stream)
+    root_logger.addHandler(handler)
+
+    def td():
+        root_logger.removeHandler(handler)
+
+    request.addfinalizer(td)
+    return log_stream
+
+
 class TestTerraform(object):
     def teardown_method(self, method):
         """ teardown any state that was previously setup with a setup_method
@@ -82,12 +102,14 @@ class TestTerraform(object):
             assert s in result
 
     @pytest.mark.parametrize(*CMD_CASES)
-    def test_cmd(self, method, expected_output, expected_ret_code):
-        os.environ['AWS_DEFAULT_REGION'] = "us-west-1"
+    def test_cmd(self, method, expected_output, expected_ret_code, expected_logs, string_logger):
         tf = Terraform(working_dir=current_path)
         ret, out, err = method(tf)
+        logs = str(string_logger.getvalue())
+        logs = logs.replace('\n', '')
         assert expected_output in out
         assert expected_ret_code == ret
+        assert expected_logs in logs
 
     @pytest.mark.parametrize(
         ("folder", "variables", "var_files", "expected_output", "options"),
@@ -172,9 +194,9 @@ class TestTerraform(object):
         ret, out, err = tf.fmt(diff=True)
         assert ret == 0
 
-    def test_import(self):
-        os.environ['AWS_DEFAULT_REGION'] = "us-west-1"
-        tf = Terraform(working_dir=current_path)
-        ret, out, err = tf.import_cmd('aws_instance.foo', 'i-abc1234', no_color=IsFlagged)
-        assert 'Import complete!' in out
-        assert 1 == ret
+    def test_import(self, string_logger):
+        tf = Terraform(working_dir=os.path.join(current_path, 'import_test'))
+        tf.import_cmd('aws_instance.foo', 'i-abc1234', no_color=IsFlagged)
+        logs = string_logger.getvalue()
+        print(logs)
+        assert 'command: terraform import -no-color aws_instance.foo i-abc1234' in logs
diff --git a/test/vars_require_input/main.tf b/test/vars_require_input/main.tf
index 43e8e65..d3eaed3 100644
--- a/test/vars_require_input/main.tf
+++ b/test/vars_require_input/main.tf
@@ -10,11 +10,11 @@ variable "list" {
 
 variable "map" {
   default = {}
-  type = "map"
+  type    = "map"
 }
 
 resource "aws_instance" "bar" {
   foo = "${var.ami}"
   bar = "${join(",", var.list)}"
   baz = "${join(",", keys(var.map))}"
-}
\ No newline at end of file
+}