From 94c68e4e1940eb0c420cf96be4151df13e5fab60 Mon Sep 17 00:00:00 2001
From: Srinivas Lade <srinulade1@gmail.com>
Date: Sun, 5 Mar 2023 15:27:44 -0500
Subject: [PATCH] Test for Multiple Categories

---
 tests/test_conda_lock.py | 50 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/tests/test_conda_lock.py b/tests/test_conda_lock.py
index 971e4c7b..db0de206 100644
--- a/tests/test_conda_lock.py
+++ b/tests/test_conda_lock.py
@@ -36,6 +36,7 @@
     _add_auth_to_line,
     _add_auth_to_lockfile,
     _extract_domain,
+    _solve_for_arch,
     _strip_auth_from_line,
     _strip_auth_from_lockfile,
     create_lockfile_from_spec,
@@ -592,6 +593,7 @@ def test_choose_wheel() -> None:
         platform="linux-64",
     )
     assert len(solution) == 1
+    assert solution["fastavro"].categories == {"main"}
     assert solution["fastavro"].hash == HashModel(
         sha256="a111a384a786b7f1fd6a8a8307da07ccf4d4c425084e2d61bae33ecfb60de405"
     )
@@ -1578,6 +1580,54 @@ def test_run_lock_with_input_hash_check(
     assert "Spec hash already locked for" in output.err
 
 
+def test_solve_arch_multiple_categories():
+    _conda_exe = determine_conda_executable(None, mamba=False, micromamba=False)
+    vpr = default_virtual_package_repodata()
+    channels = [Channel.from_string("conda-forge")]
+
+    with vpr, tempfile.NamedTemporaryFile(dir=".") as tf:
+        spec = LockSpecification(
+            dependencies={
+                "linux-64": [
+                    VersionedDependency(
+                        name="python",
+                        version="=3.10.9",
+                        manager="conda",
+                        category="main",
+                        extras=[],
+                    ),
+                    VersionedDependency(
+                        name="pandas",
+                        version="=1.5.3",
+                        manager="conda",
+                        category="test",
+                        extras=[],
+                    ),
+                    VersionedDependency(
+                        name="pyarrow",
+                        version="=9.0.0",
+                        manager="conda",
+                        category="dev",
+                        extras=[],
+                    ),
+                ],
+            },
+            channels=channels,
+            # NB: this file must exist for relative path resolution to work
+            # in create_lockfile_from_spec
+            sources=[Path(tf.name)],
+            virtual_package_repo=vpr,
+        )
+
+        locked_deps = _solve_for_arch(_conda_exe, spec, "linux-64", channels)
+        python_deps = [dep for dep in locked_deps if dep.name == "python"]
+        assert len(python_deps) == 1
+        assert python_deps[0].categories == {"main", "test", "dev"}
+        numpy_deps = [dep for dep in locked_deps if dep.name == "numpy"]
+        assert len(numpy_deps) == 1
+        assert numpy_deps[0].categories == {"test", "dev"}
+
+
 @pytest.mark.parametrize(
     "package,version,url_pattern",
     [