From 47431ddd9351ed146cc94b1ee81a1172306c8c1c Mon Sep 17 00:00:00 2001
From: Alexey Lavrenuke <direvius@gmail.com>
Date: Wed, 11 Feb 2015 17:19:44 +0300
Subject: [PATCH] Add enum_ammo parameter

---
 yandextank/plugins/Phantom/PhantomUtils.py |  3 ++-
 yandextank/plugins/Phantom/plugin.py       | 10 +++++++++-
 yandextank/stepper/config.py               |  3 ++-
 yandextank/stepper/main.py                 | 11 +++++++++--
 4 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/yandextank/plugins/Phantom/PhantomUtils.py b/yandextank/plugins/Phantom/PhantomUtils.py
index 59d1a17..5baa828 100644
--- a/yandextank/plugins/Phantom/PhantomUtils.py
+++ b/yandextank/plugins/Phantom/PhantomUtils.py
@@ -33,6 +33,7 @@ class PhantomConfig:
         self.phantom_modules_path = None
         self.threads = None
         self.additional_libs = None
+        self.enum_ammo = False
 
     def get_option(self, opt_name, default=None):
         """ get option wrapper """
@@ -41,7 +42,7 @@ class PhantomConfig:
     @staticmethod
     def get_available_options():
         opts = ["threads", "phantom_modules_path",
-                "additional_libs", "writelog", ]
+                "additional_libs", "writelog", "enum_ammo", ]
         opts += StreamConfig.get_available_options()
         return opts
 
diff --git a/yandextank/plugins/Phantom/plugin.py b/yandextank/plugins/Phantom/plugin.py
index f5e3064..a3c9709 100644
--- a/yandextank/plugins/Phantom/plugin.py
+++ b/yandextank/plugins/Phantom/plugin.py
@@ -7,6 +7,7 @@ import subprocess
 import sys
 import time
 import datetime
+import string
 
 from yandextank.plugins import ConsoleScreen
 from yandextank.plugins.Aggregator import \
@@ -44,6 +45,8 @@ class PhantomPlugin(AbstractPlugin, AggregateResultListener):
         self.cached_info = None
         self.phantom_stderr = None
 
+        self.enum_ammo = False
+
     @staticmethod
     def get_key():
         return __file__
@@ -60,6 +63,7 @@ class PhantomPlugin(AbstractPlugin, AggregateResultListener):
         self.eta_file = self.get_option("eta_file", '')
         self.core.add_artifact_file(self.eta_file)
         self.phantom_path = self.get_option("phantom_path", 'phantom')
+        self.enum_ammo = self.get_option("enum_ammo", False)
         self.buffered_seconds = int(
             self.get_option("buffered_seconds", self.buffered_seconds))
 
@@ -383,6 +387,7 @@ class PhantomReader(AbstractReader):
         self.last_sample_time = 0
         self.read_lines_count = 0
         self.buffered_seconds = 3
+        self.enum_ammo = self.phantom.enum_ammo
 
     def check_open_files(self):
         info = self.phantom.get_info()
@@ -493,7 +498,10 @@ class PhantomReader(AbstractReader):
             # bytes:     sent    received
             # connect    send    latency    receive
             #        accuracy
-            data_item = (data[1], active, rt_real / 1000, data[11], data[10],
+            marker = data[1]
+            if self.enum_ammo:
+                marker = string.rsplit(marker, "#", 1)[0]
+            data_item = (marker, active, rt_real / 1000, data[11], data[10],
                          int(data[8]), int(data[9]),
                          int(data[3]) / 1000, int(data[4]) / 1000,
                          int(data[5]) / 1000, int(data[6]) / 1000,
diff --git a/yandextank/stepper/config.py b/yandextank/stepper/config.py
index fea47ac..915679d 100644
--- a/yandextank/stepper/config.py
+++ b/yandextank/stepper/config.py
@@ -22,6 +22,7 @@ class ComponentFactory():
         uris=None,
         headers=None,
         autocases=None,
+        enum_ammo=False,
         ammo_type='phantom',
         chosen_cases=[],
     ):
@@ -47,7 +48,7 @@ class ComponentFactory():
         if self.uris and loop_limit:
             info.status.ammo_limit = len(self.uris) * loop_limit
         self.headers = headers
-        self.marker = get_marker(autocases)
+        self.marker = get_marker(autocases, enum_ammo)
         self.chosen_cases = chosen_cases
 
     def get_load_plan(self):
diff --git a/yandextank/stepper/main.py b/yandextank/stepper/main.py
index 9107c89..8227a6d 100644
--- a/yandextank/stepper/main.py
+++ b/yandextank/stepper/main.py
@@ -100,6 +100,7 @@ class StepperWrapper(object):
         self.uris = []
         self.headers = []
         self.autocases = 0
+        self.enum_ammo = False
         self.use_caching = True
         self.force_stepping = None
         self.ammo_type = 'phantom'
@@ -126,8 +127,12 @@ class StepperWrapper(object):
         opts = [StepperWrapper.OPTION_AMMOFILE, StepperWrapper.OPTION_LOOP,
                 StepperWrapper.OPTION_SCHEDULE, StepperWrapper.OPTION_STPD]
         opts += ["instances_schedule", "uris",
-                 "headers", "header_http", "autocases", "ammo_type", "ammo_limit"]
-        opts += ["use_caching", "cache_dir", "force_stepping", "file_cache", "chosen_cases"]
+                 "headers", "header_http",
+                 "autocases", "enum_ammo",
+                 "ammo_type", "ammo_limit"]
+        opts += ["use_caching", "cache_dir",
+                 "force_stepping", "file_cache",
+                 "chosen_cases"]
         return opts
 
     def read_config(self):
@@ -159,6 +164,7 @@ class StepperWrapper(object):
         self.headers = rx.findall(self.get_option("headers", ''))
         self.http_ver = self.get_option("header_http", self.http_ver)
         self.autocases = self.get_option("autocases", '0')
+        self.enum_ammo = self.get_option("enum_ammo", False)
         self.use_caching = int(self.get_option("use_caching", '1'))
 
         self.file_cache = int(self.get_option('file_cache', '8192'))
@@ -285,6 +291,7 @@ class StepperWrapper(object):
             uris=self.uris,
             headers=[header.strip('[]') for header in self.headers],
             autocases=self.autocases,
+            enum_ammo=self.enum_ammo,
             ammo_type=self.ammo_type,
             chosen_cases=self.chosen_cases,
         )