From d498cb3e0fad7783866912637365951d77c932e1 Mon Sep 17 00:00:00 2001
From: Bruno Oliveira <nicoddemus@gmail.com>
Date: Thu, 4 Aug 2016 20:50:54 -0300
Subject: [PATCH] Add env variables to identify workers

---
 README.rst             | 11 +++++++++++
 testing/test_remote.py | 20 +++++++++++++++++++-
 xdist/remote.py        |  2 ++
 xdist/slavemanage.py   |  3 ++-
 4 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/README.rst b/README.rst
index 29fa5dd..7aab2ce 100644
--- a/README.rst
+++ b/README.rst
@@ -187,6 +187,8 @@ at once.   The specifications strings use the `xspec syntax`_.
 
 Identifying the worker process during a test
 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
+
+
 If you need to determine the identity of a worker process in
 a test or fixture, you may use the ``worker_id`` fixture to do so:
 
@@ -200,6 +202,15 @@ a test or fixture, you may use the ``worker_id`` fixture to do so:
 When ``xdist`` is disabled (running with ``-n0`` for example), then
 ``worker_id`` will return ``"master"``.
 
+Additionally, worker processes have the following environment variables
+defined:
+
+* ``PYTEST_XDIST_WORKER``: the name of the worker, e.g., ``"gw2"``.
+* ``PYTEST_XDIST_WORKER_COUNT``: the total number of workers in this session,
+  e.g., ``"4"`` when ``-n 4`` is given in the command-line.
+
+*New in version 1.15.*
+
 Specifying test exec environments in an ini file
 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 
diff --git a/testing/test_remote.py b/testing/test_remote.py
index 207b926..727aeca 100644
--- a/testing/test_remote.py
+++ b/testing/test_remote.py
@@ -40,7 +40,12 @@ class SlaveSetup:
         self.gateway = execnet.makegateway()
         self.config = config = self.testdir.parseconfigure()
         putevent = self.use_callback and self.events.put or None
-        self.slp = SlaveController(None, self.gateway, config, putevent)
+
+        class DummyMananger:
+            specs = [0, 1]
+
+        self.slp = SlaveController(DummyMananger, self.gateway, config,
+                                   putevent)
         self.request.addfinalizer(self.slp.ensure_teardown)
         self.slp.setup()
 
@@ -178,6 +183,8 @@ class TestSlaveInteractor:
         ev = slave.popevent("slavefinished")
         assert 'slaveoutput' in ev.kwargs
 
+    @pytest.mark.skipif(pytest.__version__ >= '3.0',
+                        reason='skip at module level illegal in pytest 3.0')
     def test_remote_collect_skip(self, slave):
         slave.testdir.makepyfile("""
             import py
@@ -255,3 +262,14 @@ class TestSlaveInteractor:
             ("pytest_pycollect_makeitem", "name == 'test_func'"),
             ("pytest_collectreport", "report.collector.fspath == bbb"),
         ])
+
+
+def test_remote_env_vars(testdir):
+    testdir.makepyfile('''
+        import os
+        def test():
+            assert os.environ['PYTEST_XDIST_WORKER'] in ('gw0', 'gw1')
+            assert os.environ['PYTEST_XDIST_WORKER_COUNT'] == '2'
+    ''')
+    result = testdir.runpytest('-n2', '--max-slave-restart=0')
+    assert result.ret == 0
diff --git a/xdist/remote.py b/xdist/remote.py
index 226262a..40bbeac 100644
--- a/xdist/remote.py
+++ b/xdist/remote.py
@@ -148,6 +148,8 @@ if __name__ == '__channelexec__':
     os.environ['PYTHONPATH'] = (
         importpath + os.pathsep +
         os.environ.get('PYTHONPATH', ''))
+    os.environ['PYTEST_XDIST_WORKER'] = slaveinput['slaveid']
+    os.environ['PYTEST_XDIST_WORKER_COUNT'] = str(slaveinput['slavecount'])
     # os.environ['PYTHONPATH'] = importpath
     import py
     config = remote_initconfig(option_dict, args)
diff --git a/xdist/slavemanage.py b/xdist/slavemanage.py
index 7632137..6112524 100644
--- a/xdist/slavemanage.py
+++ b/xdist/slavemanage.py
@@ -205,7 +205,8 @@ class SlaveController(object):
         self.putevent = putevent
         self.gateway = gateway
         self.config = config
-        self.slaveinput = {'slaveid': gateway.id}
+        self.slaveinput = {'slaveid': gateway.id,
+                           'slavecount': len(nodemanager.specs)}
         self._down = False
         self._shutdown_sent = False
         self.log = py.log.Producer("slavectl-%s" % gateway.id)