Repository: PrefectHQ/prefect

Issue:
Prefect context manager overwrites global contexts only once
## Description
<!-- A clear description of the bug -->

The [document](https://docs.prefect.io/core/concepts/execution.html#modifying-context-at-runtime) states we can modify global context keys by using a context manager:

> ### Modifying context at runtime
> Modifying context, even globally set context keys, at specific times is possible using a provided context manager:

However, in the attached script, the context key `foo` is overwritten only once.

## Expected Behavior
<!-- What did you expect to happen instead? -->

Modified context values do not change during a flow run.

## Reproduction
<!-- A minimal example that exhibits the behavior. -->

### context-test.py
```python
import prefect
from prefect import task, Flow

@task
def print_contexts():
    print("foo = " + prefect.context.foo)

with Flow('test') as flow:
    t1 = print_contexts()
    t2 = print_contexts()
    t1.set_downstream(t2)

with prefect.context(foo='foo'):
    flow.run()
```

### ~/.prefect/config.toml
```
[context]
foo = "global"
```

### Execution result
```
% python context-test.py
[2020-12-14 16:12:46+0900] INFO - prefect.FlowRunner | Beginning Flow run for 'test'
[2020-12-14 16:12:46+0900] INFO - prefect.TaskRunner | Task 'print_contexts': Starting task run...
foo = foo
[2020-12-14 16:12:46+0900] INFO - prefect.TaskRunner | Task 'print_contexts': Finished task run for task with final state: 'Success'
[2020-12-14 16:12:46+0900] INFO - prefect.TaskRunner | Task 'print_contexts': Starting task run...
foo = global  <-- UNEXPECTED!!
[2020-12-14 16:12:46+0900] INFO - prefect.TaskRunner | Task 'print_contexts': Finished task run for task with final state: 'Success'
[2020-12-14 16:12:46+0900] INFO - prefect.FlowRunner | Flow run SUCCESS: all reference tasks succeeded
```

## Environment
<!-- Any additional information about your environment

Optionally run `prefect diagnostics` from the command line and paste the information here. -->

```
% prefect diagnostics
{
  "config_overrides": {
    "context": {
      "foo": true
    }
  },
  "env_vars": [],
  "system_information": {
    "platform": "macOS-10.16-x86_64-i386-64bit",
    "prefect_backend": "cloud",
    "prefect_version": "0.13.19",
    "python_version": "3.8.5"
  }
}
```

Discussion:

Patch:
diff --git a/changes/pr3923.yaml b/changes/pr3923.yaml
deleted file mode 100644
index a17d5f3adb..0000000000
--- a/changes/pr3923.yaml
+++ /dev/null
@@ -1,2 +0,0 @@
-fix:
-  - "Support storing multiple local flows with the same name when using `Local` storage - [#3923](https://github.com/PrefectHQ/prefect/pull/3923)"
diff --git a/changes/pr3924.yaml b/changes/pr3924.yaml
new file mode 100644
index 0000000000..de3d6b199d
--- /dev/null
+++ b/changes/pr3924.yaml
@@ -0,0 +1,2 @@
+fix:
+  - "Fix bug in `prefect.context` contextmanager that resulted in context fields reverting to their initially configured values - [#3924](https://github.com/PrefectHQ/prefect/pull/3924)"
diff --git a/src/prefect/storage/local.py b/src/prefect/storage/local.py
index 25f6c3b620..6f9ef87b44 100644
--- a/src/prefect/storage/local.py
+++ b/src/prefect/storage/local.py
@@ -2,8 +2,6 @@ import os
 import socket
 from typing import TYPE_CHECKING, Any, Dict, List
 
-import cloudpickle
-import pendulum
 from slugify import slugify
 
 import prefect
@@ -101,8 +99,7 @@ class Local(Storage):
                 if self.stored_as_script:
                     return extract_flow_from_file(file_path=flow_location)
                 else:
-                    with open(flow_location, "rb") as f:
-                        return cloudpickle.load(f)
+                    return prefect.core.flow.Flow.load(flow_location)
             # otherwise the path is given in the module format
             else:
                 return extract_flow_from_module(module_str=flow_location)
@@ -141,13 +138,9 @@ class Local(Storage):
                 flow_location = self.path
             else:
                 flow_location = os.path.join(
-                    self.directory,
-                    slugify(flow.name),
-                    slugify(pendulum.now("utc").isoformat()),
+                    self.directory, "{}.prefect".format(slugify(flow.name))
                 )
-            os.makedirs(os.path.dirname(flow_location), exist_ok=True)
-            with open(flow_location, "wb") as f:
-                cloudpickle.dump(flow, f)
+            flow_location = flow.save(flow_location)
 
         self.flows[flow.name] = flow_location
         self._flows[flow.name] = flow
diff --git a/src/prefect/utilities/context.py b/src/prefect/utilities/context.py
index d2bde4720e..27f4544bc6 100644
--- a/src/prefect/utilities/context.py
+++ b/src/prefect/utilities/context.py
@@ -82,10 +82,14 @@ class Context(DotDict, threading.local):
     """
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        super().__init__(*args, **kwargs)
-        if "context" in config:
-            self.update(config.context)
-        self["config"] = merge_dicts(config, self.get("config", {}))  # order matters
+        init = {}
+        # Initialize with config context
+        init.update(config.get("context", {}))
+        # Overwrite with explicit args
+        init.update(dict(*args, **kwargs))
+        # Merge in config (with explicit args overwriting)
+        init["config"] = merge_dicts(config, init.get("config", {}))
+        super().__init__(init)
 
     def __getstate__(self) -> None:
         """
@@ -111,7 +115,8 @@ class Context(DotDict, threading.local):
             with prefect.context(dict(a=1, b=2), c=3):
                 print(prefect.context.a) # 1
         """
-        previous_context = self.copy()
+        # Avoid creating new `Context` object, copy as `dict` instead.
+        previous_context = self.__dict__.copy()
         try:
             new_context = dict(*args, **kwargs)
             if "config" in new_context:


Test patch:
diff --git a/tests/storage/test_local_storage.py b/tests/storage/test_local_storage.py
index ea4173fc3b..a9381ac1f3 100644
--- a/tests/storage/test_local_storage.py
+++ b/tests/storage/test_local_storage.py
@@ -1,5 +1,6 @@
 import os
 import socket
+import tempfile
 
 import cloudpickle
 import pytest
@@ -40,22 +41,21 @@ def test_create_local_storage_without_validation():
     assert storage.result.dir == storage.directory
 
 
-def test_add_flow_to_storage(tmpdir):
-    storage = Local(directory=str(tmpdir))
-    f = Flow("test")
-    assert f.name not in storage
+def test_add_flow_to_storage():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        storage = Local(directory=tmpdir)
+        f = Flow("test")
+        assert f.name not in storage
 
-    res = storage.add_flow(f)
+        res = storage.add_flow(f)
+        assert res.endswith("test.prefect")
+        assert f.name in storage
 
-    flow_dir = os.path.join(tmpdir, "test")
-    assert os.path.exists(flow_dir)
-    assert len(os.listdir(flow_dir)) == 1
-    assert res.startswith(flow_dir)
+        with open(os.path.join(tmpdir, "test.prefect"), "rb") as f:
+            wat = f.read()
 
-    assert f.name in storage
-
-    f2 = storage.get_flow(res)
-    assert f2.name == "test"
+    assert isinstance(wat, bytes)
+    assert cloudpickle.loads(wat).name == "test"
 
 
 def test_add_flow_file_to_storage(tmpdir):
@@ -79,13 +79,14 @@ def test_add_flow_file_to_storage(tmpdir):
     assert f.name in storage
 
 
-def test_add_flow_raises_if_name_conflict(tmpdir):
-    storage = Local(directory=str(tmpdir))
-    f = Flow("test")
-    res = storage.add_flow(f)
-    g = Flow("test")
-    with pytest.raises(ValueError, match='name "test"'):
-        storage.add_flow(g)
+def test_add_flow_raises_if_name_conflict():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        storage = Local(directory=tmpdir)
+        f = Flow("test")
+        res = storage.add_flow(f)
+        g = Flow("test")
+        with pytest.raises(ValueError, match='name "test"'):
+            storage.add_flow(g)
 
 
 def test_get_env_runner_raises():
@@ -101,12 +102,13 @@ def test_get_flow_raises_if_flow_not_present():
         s.get_flow("test")
 
 
-def test_get_flow_returns_flow(tmpdir):
-    storage = Local(directory=str(tmpdir))
-    f = Flow("test")
-    loc = storage.add_flow(f)
-    runner = storage.get_flow(loc)
-    assert runner == f
+def test_get_flow_returns_flow():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        storage = Local(directory=tmpdir)
+        f = Flow("test")
+        loc = storage.add_flow(f)
+        runner = storage.get_flow(loc)
+        assert runner == f
 
 
 def test_get_flow_from_file_returns_flow(tmpdir):
@@ -125,73 +127,81 @@ def test_get_flow_from_file_returns_flow(tmpdir):
     assert flow.run()
 
 
-def test_containment(tmpdir):
-    s = Local(directory=str(tmpdir))
-    f = Flow("test")
-    s.add_flow(f)
+def test_containment():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        s = Local(directory=tmpdir)
+        f = Flow("test")
+        s.add_flow(f)
 
-    assert True not in s
-    assert f not in s
-    assert "test" in s
-    assert Flow("other") not in s
-    assert "other" not in s
+        assert True not in s
+        assert f not in s
+        assert "test" in s
+        assert Flow("other") not in s
+        assert "other" not in s
 
 
-def test_build_returns_self(tmpdir):
-    s = Local(directory=str(tmpdir))
-    assert s.build() is s
+def test_build_returns_self():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        s = Local(directory=tmpdir)
+        assert s.build() is s
 
-    f = Flow("test")
-    s.add_flow(f)
-    assert s.build() is s
+        f = Flow("test")
+        s.add_flow(f)
+        assert s.build() is s
 
 
-def test_multiple_flows_in_storage(tmpdir):
-    s = Local(directory=str(tmpdir))
-    f = Flow("test")
-    g = Flow("other")
-    z = Flow("not")
-    f_loc = s.add_flow(f)
-    g_loc = s.add_flow(g)
+def test_multiple_flows_in_storage():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        s = Local(directory=tmpdir)
+        f = Flow("test")
+        g = Flow("other")
+        z = Flow("not")
+        f_loc = s.add_flow(f)
+        g_loc = s.add_flow(g)
 
-    assert "test" in s
-    assert "other" in s
-    assert "not" not in s
+        assert "test" in s
+        assert "other" in s
+        assert "not" not in s
 
-    assert s.get_flow(f_loc) == f
-    assert s.get_flow(g_loc) == g
+        assert s.get_flow(f_loc) == f
+        assert s.get_flow(g_loc) == g
 
-    assert s.flows["test"] == f_loc
-    assert s.flows["other"] == g_loc
+        assert s.flows["test"] == f_loc
+        assert s.flows["other"] == g_loc
 
 
-def test_add_flow_with_weird_name_is_cleaned(tmpdir):
-    s = Local(directory=str(tmpdir))
-    flow = Flow("WELL what do you know?!~? looks like a test!!!!")
-    loc = s.add_flow(flow)
-    assert "?" not in loc
-    assert "!" not in loc
-    assert " " not in loc
-    assert "~" not in loc
+def test_add_flow_with_weird_name_is_cleaned():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        s = Local(directory=tmpdir)
+        flow = Flow("WELL what do you know?!~? looks like a test!!!!")
+        loc = s.add_flow(flow)
+        assert "?" not in loc
+        assert "!" not in loc
+        assert " " not in loc
+        assert "~" not in loc
 
 
-def test_build_healthchecks(tmpdir):
-    s = Local(directory=str(tmpdir))
-    flow = Flow("TestFlow")
-    s.add_flow(flow)
-    assert s.build()
+def test_build_healthchecks():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        s = Local(directory=tmpdir)
+        flow = Flow("TestFlow")
+        s.add_flow(flow)
+        assert s.build()
 
 
-def test_build_healthcheck_returns_on_no_flows(tmpdir):
-    s = Local(directory=str(tmpdir))
-    assert s.build()
+def test_build_healthcheck_returns_on_no_flows():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        s = Local(directory=tmpdir)
+        assert s.build()
 
 
-def test_labels_includes_hostname(tmpdir):
-    s = Local(directory=str(tmpdir))
-    assert socket.gethostname() in s.labels
+def test_labels_includes_hostname():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        s = Local(directory=tmpdir)
+        assert socket.gethostname() in s.labels
 
 
-def test_opt_out_of_hostname_label(tmpdir):
-    s = Local(directory=str(tmpdir), add_default_labels=False)
-    assert socket.gethostname() not in s.labels
+def test_opt_out_of_hostname_label():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        s = Local(directory=tmpdir, add_default_labels=False)
+        assert socket.gethostname() not in s.labels
diff --git a/tests/test_context.py b/tests/test_context.py
index 79f7a2c0f6..b17a2a3fc0 100644
--- a/tests/test_context.py
+++ b/tests/test_context.py
@@ -120,11 +120,13 @@ def test_modify_context_by_calling_update_inside_contextmanager():
 
 def test_context_loads_values_from_config(monkeypatch):
     subsection = Config(password="1234")
-    config = Config(context=Config(subsection=subsection, my_key="my_value"))
+    config = Config(context=Config(subsection=subsection, key1="val1", key2="val2"))
     monkeypatch.setattr(prefect.utilities.context, "config", config)
-    fresh_context = Context()
+
+    fresh_context = Context(key2="new")
     assert "subsection" in fresh_context
-    assert fresh_context.my_key == "my_value"
+    assert fresh_context.key1 == "val1"
+    assert fresh_context.key2 == "new"  # overridden by constructor
     assert fresh_context.subsection == subsection
 
 
@@ -137,6 +139,27 @@ def test_context_loads_secrets_from_config(monkeypatch):
     assert fresh_context.secrets == secrets_dict
 
 
+def test_context_contextmanager_prioritizes_new_keys_even_on_context_exit(monkeypatch):
+    """Previously exiting a context block would reload from the config,
+    overwriting any explicitly set values in a nested context. This was due to
+    the `Context` constructor being implicitly called when stashing the old
+    context, and the constructor prioritizing `config.context` over explicit
+    values."""
+    config = Config(context=Config(my_key="fizz"))
+    monkeypatch.setattr(prefect.utilities.context, "config", config)
+
+    context = Context()
+    assert context.my_key == "fizz"
+
+    with context(my_key="buzz"):
+        assert context.my_key == "buzz"
+        with context({"config": {"logging": {"log_to_cloud": "FOO"}}}):
+            assert context.config.logging.log_to_cloud == "FOO"
+            assert context.my_key == "buzz"
+        assert context.my_key == "buzz"
+    assert context.my_key == "fizz"
+
+
 def test_context_contextmanager_prioritizes_new_config_keys():
     with prefect.context({"config": {"logging": {"log_to_cloud": "FOO"}}}):
         assert prefect.context.config.logging.log_to_cloud == "FOO"